[Rdkit-devel] [PATCH 12/17] Add highlightMap parameter to allow arbitrary coloring on atoms

[Gianluca S. <gi...@gm...>]

---
 rdkit/Chem/Draw/MolDrawing.py |   32 ++++++++++++++++++++------------
 rdkit/Chem/Draw/__init__.py   |    3 ++-
 2 files changed, 22 insertions(+), 13 deletions(-)

diff --git a/rdkit/Chem/Draw/MolDrawing.py b/rdkit/Chem/Draw/MolDrawing.py
index 0cc24b4..8498704 100644
--- a/rdkit/Chem/Draw/MolDrawing.py
+++ b/rdkit/Chem/Draw/MolDrawing.py
@@ -266,11 +266,8 @@ class MolDrawing(object):
     drawingTrans = canvasSize[0]/2,canvasSize[1]/2
     self.drawingTrans = drawingTrans
 
-  def _drawLabel(self,label,pos,font,color=None,
-                 highlightIt=False):
-    if highlightIt:
-      color = self.selectColor
-    elif not color:
+  def _drawLabel(self,label,pos,font,color=None):
+    if color is None:
       color = self.defaultColor
     x1 = pos[0]
     y1 = pos[1]
@@ -278,13 +275,16 @@ class MolDrawing(object):
     self.canvas.addCanvasText(label,(x1,y1),font,color)
     
   def AddMol(self,mol,centerIt=True,molTrans=None,drawingTrans=None,
-             highlightAtoms=[],confId=-1,flagCloseContactsDist=2
-             ignoreHs=False):
-    """
+             highlightAtoms=[],confId=-1,flagCloseContactsDist=2,
+             highlightMap=None, ignoreHs=False):
+    """Set the molecule to be drawn.
+
+    Parameters:
+      hightlightAtoms -- list of atoms to highlight (default [])
+      highlightMap -- dictionary of (atom, color) pairs (default None)
 
     Notes:
       - specifying centerIt will cause molTrans and drawingTrans to be ignored
-      
     """
     conf = mol.GetConformer(confId)
 
@@ -342,6 +342,10 @@ class MolDrawing(object):
             width=2.0*self.bondLineWidth
             color = self.selectColor
             color2 = self.selectColor
+          elif highlightMap is not None and idx in highlightMap and nbrIdx in highlightMap:
+            width=2.0*self.bondLineWidth
+            color = highlightMap[idx]
+            color2 = highlightMap[nbrIdx]
           else:
             width=self.bondLineWidth
             if self.colorBonds:
@@ -399,9 +403,13 @@ class MolDrawing(object):
           else:
             symbol = '%s%s%s'%(chg,hs,base)
 
-        color = self.elemDict.get(atom.GetAtomicNum(),(0,0,0))
-        self._drawLabel(symbol,pos,font,color=color,
-                        highlightIt=(highlightAtoms and idx in highlightAtoms))
+        if highlightMap and idx in highlightMap:
+          color = highlightMap[idx]
+        elif highlightAtoms and idx in highlightAtoms:
+          color = self.selectColor
+        else:
+          color = self.elemDict.get(atom.GetAtomicNum(),(0,0,0))
+        self._drawLabel(symbol, pos, font, color=color)
 
     if flagCloseContactsDist>0:
       tol = flagCloseContactsDist*flagCloseContactsDist
diff --git a/rdkit/Chem/Draw/__init__.py b/rdkit/Chem/Draw/__init__.py
index 19b9a7d..e1ec134 100644
--- a/rdkit/Chem/Draw/__init__.py
+++ b/rdkit/Chem/Draw/__init__.py
@@ -11,10 +11,11 @@ def MolToImage(mol, size=(300,300), kekulize=True, wedgeBonds=True, **kwargs):
   """ returns a PIL image containing a drawing of the molecule
 
     Keyword arguments:
-    kekulize -- run kekulization routine on input `mol`
+    kekulize -- run kekulization routine on input `mol` (default True)
     size -- final image size, in pixel (default (300,300))
     wedgeBonds -- draw wedge (stereo) bonds (default True)
     highlightAtoms -- list of atoms to highlight (default [])
+    highlightMap -- dictionary of (atom, color) pairs (default None)
   """
   if not mol:
     raise ValueError,'Null molecule provided'
-- 
1.7.4