--- a/files/test.sh	Tue Sep 06 23:59:06 2022 +0300
+++ b/files/test.sh	Wed Sep 07 14:23:52 2022 +0300
@@ -1,6 +1,6 @@
 #!/bin/bash
 
-python ../medchem_fragment_splitter.py -l DB08865.smi -w . -p -x -a
+python3 ../medchem_fragment_splitter.py -l DB08865.smi -w . -p -x -a
 SVG_DB08865=$(diff DB08865.smi.svg output_DB08865.smi.svg | wc -c)
 if [[ ${SVG_DB08865} -eq '0' ]]; then
 	echo "Draw SVG Test 1 OK!"
@@ -31,15 +31,15 @@
 fi
 
 # clean up
-rm DB08865.smi.svg
-rm DB08865.smi.png
+#rm DB08865.smi.svg
+#rm DB08865.smi.png
 rm ligand_1_C2H4.pdb
 rm ligand_2_C8H8N4O.pdb
 rm ligand_3_C5H11N.pdb
 rm ligand_4_C6H3Cl2F.pdb
 
 # Test 4-6
-python ../medchem_fragment_splitter.py -l DB00176.smi -w . -p -x -a
+python3 ../medchem_fragment_splitter.py -l DB00176.smi -w . -p -x -a
 SVG_DB00176=$(diff DB00176.smi.svg output_DB00176.smi.svg | wc -c)
 if [[ ${SVG_DB00176} -eq '0' ]]; then
 	echo "Draw SVG Test 4 OK!"
@@ -72,10 +72,10 @@
 fi
 
 # clean up
-rm DB00176.smi.svg
-rm DB00176.smi.png
+#rm DB00176.smi.svg
+#rm DB00176.smi.png
 rm ligand_1_CH4O.pdb
 rm ligand_2_C4H10.pdb
 rm ligand_3_C7H7NO.pdb
 rm ligand_4_C2H7N.pdb
-rm ligand_5_CHF3.pdb
\ No newline at end of file
+rm ligand_5_CHF3.pdb

--- a/medchem_fragment_splitter.py	Tue Sep 06 23:59:06 2022 +0300
+++ b/medchem_fragment_splitter.py	Wed Sep 07 14:23:52 2022 +0300
@@ -152,19 +152,17 @@
     for bbond in break_bonds:
         bondsList.append(bbond)
         colorb[bbond] = (1.0, 0.0, 0.0)
-    rdMolDraw2D.PrepareAndDrawMolecule(imageStructure, molecule, highlightBonds=bondsList, highlightBondColors=colorb)
-    rdMolDraw2D.PrepareAndDrawMolecule(imagePNGStructure, molecule, highlightBonds=bondsList, highlightBondColors=colorb)
+    highlightAtoms = list()
+    col_atoms = {}
+    atom_radii = {}
+    red_c = 0.6
+    green_c = 0.84
+    blue_c = 0.6
+    current_color = (red_c, green_c, blue_c)
+
     for index in range(len(fragments)):
-        red_c = 0.6
-        green_c = 0.84
-        blue_c = 0.6
-
+        highlightAtoms.extend(ifragments[index])
         numAtomsInFragment = len(ifragments[index])
-        col_atoms = {}
-        current_color = (red_c, green_c, blue_c)
-        atom_radii = {}
-        colorb = {}
-        bondsList = list()
         for atom1 in range(numAtomsInFragment):
             col_atoms[ifragments[index][atom1]] = current_color
             atom_radii[ifragments[index][atom1]] = 0.45
@@ -177,14 +175,18 @@
                         bondsList.append(bond_id)
                 except Exception:
                     pass
-        rdMolDraw2D.PrepareAndDrawMolecule(imageStructure, molecule, highlightAtoms=list(ifragments[index]),
+    
+    # draw all at once
+    rdMolDraw2D.PrepareAndDrawMolecule(imageStructure, molecule, highlightAtoms=highlightAtoms,
                                            highlightAtomRadii=atom_radii,
                                            highlightBonds=bondsList,
                                            highlightAtomColors=col_atoms, highlightBondColors=colorb)
-        rdMolDraw2D.PrepareAndDrawMolecule(imagePNGStructure, molecule, highlightAtoms=list(ifragments[index]),
+    rdMolDraw2D.PrepareAndDrawMolecule(imagePNGStructure, molecule, highlightAtoms=highlightAtoms,
                                            highlightAtomRadii=atom_radii,
                                            highlightBonds=bondsList,
                                            highlightAtomColors=col_atoms, highlightBondColors=colorb)
+    imageStructure.FinishDrawing()
+    imagePNGStructure.FinishDrawing()
     svg = imageStructure.GetDrawingText().replace('svg:', '')
     print(svg, file=open(structure_output+".svg", 'w'))
     imagePNGStructure.WriteDrawingText(structure_output+".png")
@@ -264,4 +266,4 @@
 ]
 
 if __name__ == '__main__':
-    decomposition(sys.argv[1:])
\ No newline at end of file
+    decomposition(sys.argv[1:])