cnotebook 1.2.0__py3-none-any.whl → 2.1.1__py3-none-any.whl

cnotebook/helpers.py

@@ -1,7 +1,17 @@
+import logging
 import re
-from typing import Callable
+from typing import Callable, Literal, Sequence
 from openeye import oechem, oedepict
 
+log = logging.getLogger("cnotebook")
+
+
+# Type alias for highlight style
+HighlightStyle = int | Literal["overlay_default", "overlay_ball_and_stick"]
+
+# Type alias for highlight colors
+HighlightColors = oechem.OEColor | oechem.OEColorIter
+
 
 def escape_html(val):
     """
@@ -44,26 +54,148 @@ def remove_omega_conformer_id(val):
 
 def create_structure_highlighter(
         query: str | oechem.OESubSearch | oechem.OEMCSSearch | oechem.OEQMol,
-        color: oechem.OEColor = oechem.OEColor(oechem.OELightBlue),
-        style: int = oedepict.OEHighlightStyle_Stick
+        color: HighlightColors | None = None,
+        style: HighlightStyle = "overlay_default"
 ) -> Callable[[oedepict.OE2DMolDisplay], None]:
     """
-    Closure that creates a callback to highlight SMARTS patterns or MCSS results in a molecule
-    :param query: SMARTS pattern, oechem.OESubSearch, or oechem.OEMCSSearch object
-    :param color: Highlight color
-    :param style: Highlight style
-    :return: Function that highlights structures
+    Closure that creates a callback to highlight SMARTS patterns or MCSS results in a molecule.
+
+    :param query: SMARTS pattern, oechem.OESubSearch, or oechem.OEMCSSearch object.
+    :param color: Highlight color(s). Can be a single oechem.OEColor or an oechem.OEColorIter
+        (e.g., oechem.OEGetLightColors()). Defaults to oechem.OEGetLightColors().
+    :param style: Highlight style. Can be an int (OEHighlightStyle constant) or a string
+        ("overlay_default", "overlay_ball_and_stick"). Defaults to "overlay_default".
+    :returns: Callback function that highlights structures.
     """
+    # Default color
+    if color is None:
+        color = oechem.OEGetLightColors()
 
+    # Create the substructure search object
     if isinstance(query, (str, oechem.OEQMol)):
         ss = oechem.OESubSearch(query)
-
-    elif not isinstance(query, (oechem.OESubSearch, oechem.OEMCSSearch)):
+    elif isinstance(query, (oechem.OESubSearch, oechem.OEMCSSearch)):
+        ss = query
+    else:
         raise TypeError(f'Cannot create structure highlighter with object pattern of type {type(query).__name__}')
 
-    # Create the callback as a closure
-    def _structure_highlighter(disp: oedepict.OE2DMolDisplay):
-        for match in ss.Match(disp.GetMolecule(), True):
-            oedepict.OEAddHighlighting(disp, color, style, match)
+    # Determine highlighting approach based on style
+    use_overlay = isinstance(style, str) and style in ("overlay_default", "overlay_ball_and_stick")
+
+    # Check if color is compatible with overlay
+    if use_overlay and isinstance(color, oechem.OEColor):
+        log.warning(
+            "Overlay coloring is not compatible with a single oechem.OEColor. Falling back to standard highlighting")
+        use_overlay = False
+        style = oedepict.OEHighlightStyle_BallAndStick
+
+    if use_overlay:
+        # Overlay highlighting with color iterator
+        highlight = oedepict.OEHighlightOverlayByBallAndStick(color)
+
+        def _overlay_highlighter(disp: oedepict.OE2DMolDisplay):
+            oedepict.OEAddHighlightOverlay(disp, highlight, ss.Match(disp.GetMolecule(), True))
+
+        return _overlay_highlighter
+
+    else:
+        # Traditional highlighting with OEHighlightStyle int
+        # For traditional highlighting, we need a single color
+        if isinstance(color, oechem.OEColor):
+            highlight_color = color
+        else:
+            # Get first color from iterator for traditional highlighting
+            highlight_color = oechem.OELightBlue
+            for c in color:
+                highlight_color = c
+                break
+
+        def _structure_highlighter(disp: oedepict.OE2DMolDisplay):
+            for match in ss.Match(disp.GetMolecule(), True):
+                oedepict.OEAddHighlighting(disp, highlight_color, style, match)
+
+        return _structure_highlighter
+
+
+def highlight_smarts(
+    mol: oechem.OEMolBase,
+    smarts: str | Sequence[str],
+    color: oechem.OEColor | Sequence[oechem.OEColor] = oechem.OEColor(oechem.OELightBlue),
+    style: int | Sequence[int] = oedepict.OEHighlightStyle_Stick,
+    opts: oedepict.OE2DMolDisplayOptions | None = None,
+) -> oedepict.OE2DMolDisplay:
+    """Highlight SMARTS patterns in a molecule and return a display object.
+
+    :param mol: OpenEye molecule to highlight.
+    :param smarts: SMARTS pattern or sequence of SMARTS patterns to highlight.
+    :param color: Highlight color or sequence of colors. If a single color, it is
+        applied to all patterns. If a sequence, must match length of smarts.
+    :param style: Highlight style or sequence of styles. If a single style, it is
+        applied to all patterns. If a sequence, must match length of smarts.
+    :param opts: Display options. If None, default options are used.
+    :returns: OE2DMolDisplay object with highlighted substructures.
+    :raises ValueError: If color/style sequence length doesn't match smarts length.
+
+    Example::
+
+        from openeye import oechem
+        import cnotebook
+
+        mol = oechem.OEGraphMol()
+        oechem.OESmilesToMol(mol, "c1ccc2c(c1)nc(n2)N")
+
+        # Single color/style for all patterns
+        disp = cnotebook.highlight_smarts(mol, ["ncn", "c1ccccc1"])
+
+        # Different colors for each pattern
+        disp = cnotebook.highlight_smarts(
+            mol,
+            ["ncn", "c1ccccc1"],
+            color=[oechem.OEColor(oechem.OELightBlue), oechem.OEColor(oechem.OEPink)]
+        )
+    """
+    # Prepare molecule for depiction
+    oedepict.OEPrepareDepiction(mol)
+
+    # Create display options if not provided
+    if opts is None:
+        opts = oedepict.OE2DMolDisplayOptions()
+
+    # Create display object
+    disp = oedepict.OE2DMolDisplay(mol, opts)
+
+    # Normalize smarts to a list
+    if isinstance(smarts, str):
+        smarts = [smarts]
+
+    n_patterns = len(smarts)
+
+    # Normalize colors to a list
+    if isinstance(color, oechem.OEColor):
+        colors = [color] * n_patterns
+    else:
+        colors = list(color)
+        if len(colors) != n_patterns:
+            raise ValueError(
+                f"Length of color sequence ({len(colors)}) must match "
+                f"length of smarts sequence ({n_patterns})"
+            )
+
+    # Normalize styles to a list
+    if isinstance(style, int):
+        styles = [style] * n_patterns
+    else:
+        styles = list(style)
+        if len(styles) != n_patterns:
+            raise ValueError(
+                f"Length of style sequence ({len(styles)}) must match "
+                f"length of smarts sequence ({n_patterns})"
+            )
+
+    # Highlight all patterns with corresponding color and style
+    for pattern, c, s in zip(smarts, colors, styles):
+        subs = oechem.OESubSearch(pattern)
+        for match in subs.Match(mol, True):
+            oedepict.OEAddHighlighting(disp, c, s, match)
 
-    return _structure_highlighter
+    return disp

cnotebook/ipython_ext.py

@@ -1,12 +1,9 @@
 import logging
 from openeye import oechem, oedepict
 from .render import (
-    CNotebookContext,
-    pass_cnotebook_context,
     oemol_to_html,
     oedisp_to_html,
     oeimage_to_html,
-    render_molecule_grid  # Re-export for backward compatibility
 )
 
 # Only register iPython formatters if that is present

cnotebook/marimo_ext.py

@@ -11,11 +11,22 @@ from openeye import oechem, oedepict
 
 # Import oepandas for dtype checking
 try:
+    # noinspection PyUnusedImports
     import oepandas as oepd
     oepandas_available = True
 except ImportError:
     oepandas_available = False
 
+# Import oepolars for dtype checking
+try:
+    # noinspection PyUnusedImports
+    import polars as pl
+    # noinspection PyUnusedImports
+    import oepolars as oeplr
+    oepolars_available = True
+except ImportError:
+    oepolars_available = False
+
 from .context import cnotebook_context, get_series_context
 from .render import (
     oemol_to_html,
@@ -184,6 +195,48 @@ try:
     # Do the installation
     install_marimo_pandas_formatter()
 
+    def marimo_polars_formatter(df: pl.DataFrame):
+        """
+        Marimo DataFrame formatter for Polars DataFrames with molecule columns.
+        """
+        format_mapping = {}
+
+        # Check for MoleculeType / DisplayType (OEPolars specific)
+        if oepolars_available:
+            for col in df.columns:
+                dtype = df.schema[col]
+
+                if isinstance(dtype, oeplr.MoleculeType):
+                    series = df.get_column(col)
+                    metadata = series.chem.metadata if hasattr(series, 'chem') else {}
+                    ctx = get_series_context(metadata).copy()
+                    format_mapping[col] = _create_molecule_formatter(ctx)
+
+                elif isinstance(dtype, oeplr.DisplayType):
+                    series = df.get_column(col)
+                    metadata = series.chem.metadata if hasattr(series, 'chem') else {}
+                    ctx = get_series_context(metadata).copy()
+                    format_mapping[col] = _create_display_formatter(ctx)
+
+        # Return a Marimo table with our custom mapping
+        # noinspection PyProtectedMember,PyTypeChecker
+        return table(df, selection=None, format_mapping=format_mapping, pagination=True)._mime_()
+
+    def install_marimo_polars_formatter():
+        """Install the Polars DataFrame formatter if polars is available."""
+        if not oepolars_available:
+            return
+
+        # Check if we've already installed it to avoid duplicates
+        for typ, func in OPINIONATED_FORMATTERS.formatters.items():
+            if typ is pl.DataFrame and func.__name__ == "marimo_polars_formatter":
+                return  # Already installed
+
+        OPINIONATED_FORMATTERS.formatters[pl.DataFrame] = marimo_polars_formatter
+
+    if oepolars_available:
+        install_marimo_polars_formatter()
+
 except (ImportError, AttributeError) as ex:
     # Marimo not installed or API changed - skip formatter registration
     log.debug(f'Marimo formatter registration skipped: {ex}')
@@ -203,3 +256,17 @@ def _display_dataframe(self: pd.DataFrame):
     return "text/html", render_dataframe(df=self, formatters=None, col_space=None)
 
 pd.DataFrame._mime_ = _display_dataframe
+
+if oepolars_available:
+    from .polars_ext import render_polars_dataframe
+
+    def _display_polars_dataframe(self: pl.DataFrame):
+        """
+        Fallback MIME hook for Polars DataFrames in non-Marimo contexts.
+
+        In Marimo, the internal table patch handles DataFrame display.
+        This is used for static exports or other tools that check _mime_.
+        """
+        return "text/html", render_polars_dataframe(df=self, formatters=None, col_space=None)
+
+    pl.DataFrame._mime_ = _display_polars_dataframe