allow multiple substructure and text matches

Author: cbouy

molhighlighter/highlight.py

@@ -1,11 +1,28 @@
 from rdkit import Chem
 
 class Highlight:
-    def __init__(self, indices, text=None, color=None, fill_ring=False):
+    def __init__(self, indices, text=None, color=None, fill_ring=False, same_color=False):
+        """Creates a highlight from a list of atom indices
+
+        Parameters
+        ----------
+        indices : list[int] or list[list[int]]
+            Indices of atoms to highlight
+        text : str or None
+            Text to match on the molecule's label
+        color : str, tuple[float] or None
+            Color of the highlight, as a hex code or an RGB tuple. Use `None`
+            for automatic color assignment
+        fill_ring : bool
+            Fill the entire ring or simply highlight the ring atoms
+        same_color : bool
+            If a nested list is given, assign the same color to all matches
+        """
         self.indices = indices
         self.text = text
         self.color = color
         self.fill_ring = fill_ring
+        self.same_color = same_color
 
     def __repr__(self):
         ix = self.indices
@@ -15,12 +32,28 @@ def __repr__(self):
         return f'<Highlight({ix=}, {text=}, {color=}, {fill_ring=})>'
 
     @classmethod
-    def from_smarts(cls, mol, smarts, text=None, color=None, fill_ring=False):
+    def from_smarts(cls, mol, smarts, text=None, color=None, fill_ring=False,
+                    same_color=False, single_match=True):
+        """Creates a highlight from a SMARTS string
+        
+        Parameters
+        ----------
+        mol : RDKit.Chem.rdchem.Mol
+            RDKit mol on which the substructure query is performed
+        smarts : str
+            SMARTS string for the substructure query
+        single_match : bool
+            Restrict query matches to a single one
+        """
         qmol = Chem.MolFromSmarts(smarts)
-        indices = mol.GetSubstructMatch(qmol)
+        if single_match:
+            indices = mol.GetSubstructMatch(qmol)
+        else:
+            indices = mol.GetSubstructMatches(qmol)
         if not indices:
-            raise ValueError(f"Not match found for {smarts!r}")
-        return cls(indices=indices, text=text, color=color, fill_ring=fill_ring)
+            raise ValueError(f"No match found for {smarts!r}")
+        return cls(indices=indices, text=text, color=color, fill_ring=fill_ring,
+                   same_color=same_color)
 
     @property
     def color(self):

molhighlighter/molhighlighter.py

@@ -4,6 +4,7 @@
 from rdkit.Chem import AllChem, Draw
 from .utils import (Substitution, requires_config,
                     sequential_palette, get_auto_palette)
+from .highlight import Highlight
 try:
     from IPython.display import display_svg, display_html, Javascript
     from ipywidgets import ColorPicker
@@ -73,15 +74,41 @@ def configure(self, size=(-1, -1), black_font=True, fill_rings=None,
         if self.label and any(h.text is None for h in self.highlights):
             raise ValueError("Cannot use empty highlight text if label is set")
 
+        # extend highlights if indices is nested list
+        extended = []
+        for h in self.highlights:
+            if isinstance(h.indices[0], (list, tuple)) and not h.same_color:
+                ext = [Highlight(ix, h.text, h.color, h.fill_ring)
+                       for ix in h.indices]
+                extended.extend(ext)
+            else:
+                extended.append(h)
+        self.highlights = extended
+
         # automatic colors if not set
-        if all(h.color is None for h in self.highlights):
+        if any(h.color is None for h in self.highlights):
             num_hl = len(self.highlights)
             if num_hl > 5:
                 palette = get_auto_palette(num_hl)
             else:
                 palette = sequential_palette
-            for highlight, color in zip(self.highlights, palette):
-                highlight.color = color
+            for h, color in zip(self.highlights, palette):
+                h.color = h.color if h.color else color
+
+        # extend highlights if indices is nested list or text is list
+        extended = []
+        for h in self.highlights:
+            if isinstance(h.indices[0], (list, tuple)):
+                ext = [Highlight(ix, h.text, h.color, h.fill_ring)
+                       for ix in h.indices]
+                extended.extend(ext)
+            elif isinstance(h.text, (list, tuple)):
+                ext = [Highlight(h.indices, txt, h.color, h.fill_ring)
+                       for txt in h.text]
+                extended.extend(ext)
+            else:
+                extended.append(h)
+        self.highlights = extended
 
         # MolDrawOptions defaults
         opts = Draw.MolDrawOptions()
@@ -212,7 +239,7 @@ def generate_label(self):
             # find start index of text in label
             start = self._find_text(text, 0, starts)
             if start < 0:
-                warnings.warn(f"No match found in label for {highlight}")
+                warnings.warn(f"{highlight.text!r} unmatched or already found in label")
                 continue
             end = start + size
             # create substitution string