cnotebook 2.1.0__py3-none-any.whl

cnotebook/helpers.py

@@ -0,0 +1,201 @@
+import logging
+import re
+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):
+    """
+    Perform the same HTML escaping done by Pandas for displaying in Notebooks
+    :param val: Value to escape
+    :return: Escaped value (if val was a string)
+    """
+    if isinstance(val, str):
+        return val.replace("&", r"&amp;").replace("<", r"&lt;").replace(">", r"&gt;")
+    return val
+
+
+def escape_brackets(val):
+    """
+    Escapes only HTML brackets
+    :param val: Value to escape
+    :return: Escaped value (if string)
+    """
+    if isinstance(val, str):
+        return val.replace("<", r"&lt;").replace(">", r"&gt;")
+    return val
+
+
+# Remove conformer identifier from compound ID
+CONFORMER_ID_REGEX = re.compile(r'(.*?)_\d+$')
+
+
+def remove_omega_conformer_id(val):
+    """
+    Remove the conformer ID from a compound identifier
+    :param val: Value
+    :return: Processed value
+    """
+    if isinstance(val, str):
+        m = re.search(CONFORMER_ID_REGEX, val)
+        if m is not None:
+            return m.group(1)
+    return val
+
+
+def create_structure_highlighter(
+        query: str | oechem.OESubSearch | oechem.OEMCSSearch | oechem.OEQMol,
+        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(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 isinstance(query, (oechem.OESubSearch, oechem.OEMCSSearch)):
+        ss = query
+    else:
+        raise TypeError(f'Cannot create structure highlighter with object pattern of type {type(query).__name__}')
+
+    # 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 disp

cnotebook/ipython_ext.py

@@ -0,0 +1,56 @@
+import logging
+from openeye import oechem, oedepict
+from .render import (
+    oemol_to_html,
+    oedisp_to_html,
+    oeimage_to_html,
+)
+
+# Only register iPython formatters if that is present
+try:
+    # noinspection PyProtectedMember,PyPackageRequirements
+    from IPython import get_ipython
+    ipython_present = True
+except ModuleNotFoundError:
+    ipython_present = False
+
+log = logging.getLogger("cnotebook")
+
+
+########################################################################################################################
+# Register iPython formatters
+########################################################################################################################
+
+if ipython_present:
+
+    def register_ipython_formatters():
+        """
+        Register formatters for OpenEye types here that can be rendered. Calls to this function are idempotent.
+        """
+        ipython_instance = get_ipython()
+
+        if ipython_instance is not None:
+            html_formatter = ipython_instance.display_formatter.formatters['text/html']
+
+            try:
+                _ = html_formatter.lookup(oechem.OEMolBase)
+            except KeyError:
+                html_formatter.for_type(oechem.OEMolBase, oemol_to_html)
+
+            try:
+                _ = html_formatter.lookup(oedepict.OE2DMolDisplay)
+            except KeyError:
+                html_formatter.for_type(oedepict.OE2DMolDisplay, oedisp_to_html)
+
+            try:
+                _ = html_formatter.lookup(oedepict.OEImage)
+            except KeyError:
+                html_formatter.for_type(oedepict.OEImage, oeimage_to_html)
+        else:
+            log.debug("[cnotebook] iPython installed but not in use - cannot register iPython extension")
+
+else:
+
+    # iPython is not present, so we do not register formatters for OpenEye objects
+    def register_ipython_formatters():
+        pass

cnotebook/marimo_ext.py

@@ -0,0 +1,272 @@
+"""
+Marimo integration for CNotebook.
+
+This module provides MIME handlers for OpenEye objects and patches Marimo's
+internal table implementation to support molecule rendering with callbacks
+(highlighting, alignment, etc.) in Marimo's built-in DataFrame table component.
+"""
+import logging
+import pandas as pd
+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,
+    oedisp_to_html,
+    oeimage_to_html,
+    oemol_to_disp,
+    render_empty_molecule,
+    render_invalid_molecule
+)
+from .pandas_ext import render_dataframe
+
+log = logging.getLogger("cnotebook")
+
+
+########################################################################################################################
+# MIME handlers for individual OpenEye objects
+########################################################################################################################
+
+def _display_mol(self: oechem.OEMolBase):
+    ctx = cnotebook_context.get().copy()
+    # Allow user's image_format preference (SVG or PNG)
+    return "text/html", oemol_to_html(self, ctx=ctx)
+
+oechem.OEMolBase._mime_ = _display_mol
+
+
+def _display_display(self: oedepict.OE2DMolDisplay):
+    ctx = cnotebook_context.get().copy()
+    # Allow user's image_format preference (SVG or PNG)
+    return "text/html", oedisp_to_html(self, ctx=ctx)
+
+oedepict.OE2DMolDisplay._mime_ = _display_display
+
+
+def _display_image(self: oedepict.OEImage):
+    ctx = cnotebook_context.get().copy()
+    # Allow user's image_format preference (SVG or PNG)
+    return "text/html", oeimage_to_html(self, ctx=ctx)
+
+oedepict.OEImage._mime_ = _display_image
+
+
+########################################################################################################################
+# Formatter factories for mo.ui.table format_mapping
+########################################################################################################################
+
+def _create_molecule_formatter(ctx):
+    """
+    Create a formatter closure that renders molecules with callbacks applied.
+
+    :param ctx: CNotebookContext with callbacks (e.g., highlighting)
+    :return: Formatter function for use in mo.ui.table format_mapping
+    """
+    def formatter(mol):
+        if mol is None:
+            return ""
+
+        if not isinstance(mol, oechem.OEMolBase):
+            return str(mol)
+
+        # Handle invalid molecules
+        if not mol.IsValid():
+            return render_invalid_molecule(ctx=ctx)
+
+        # Handle empty molecules
+        if mol.NumAtoms() == 0:
+            return render_empty_molecule(ctx=ctx)
+
+        # Create display object
+        disp = oemol_to_disp(mol, ctx=ctx)
+
+        # Apply callbacks (highlighting, etc.)
+        if ctx.callbacks:
+            for callback in ctx.callbacks:
+                callback(disp)
+
+        # Return display object
+        return disp
+
+    return formatter
+
+
+def _create_display_formatter(ctx):
+    """
+    Create a formatter closure that renders OE2DMolDisplay objects.
+
+    :param ctx: CNotebookContext for rendering options
+    :return: Formatter function for use in mo.ui.table format_mapping
+    """
+    def formatter(disp):
+        if disp is None:
+            return ""
+
+        if not isinstance(disp, oedepict.OE2DMolDisplay):
+            return str(disp)
+
+        if not disp.IsValid():
+            return str(disp)
+
+        # Copy the display to avoid modifying the original
+        disp_copy = oedepict.OE2DMolDisplay(disp)
+
+        # Apply callbacks if any
+        if ctx.callbacks:
+            for callback in ctx.callbacks:
+                callback(disp_copy)
+
+        return disp_copy
+
+    return formatter
+
+
+########################################################################################################################
+# Marimo DataFrame formatter registration
+#
+# This registers a custom formatter with Marimo's OPINIONATED_FORMATTERS registry
+# to handle DataFrames containing molecule columns. The formatter:
+# - Detects MoleculeDtype and DisplayDtype columns
+# - Creates format_mapping entries that apply callbacks (highlighting, alignment, etc.)
+# - Returns OE2DMolDisplay objects which Marimo renders via their _mime_() method
+########################################################################################################################
+
+try:
+    import marimo as mo
+    # noinspection PyProtectedMember,PyUnusedImports
+    from marimo._output.formatting import OPINIONATED_FORMATTERS
+    # noinspection PyProtectedMember,PyUnusedImports
+    from marimo._plugins.ui._impl.table import table
+
+
+    # 1. Define the custom formatting logic
+    def marimo_pandas_formatter(df: pd.DataFrame):
+        """
+        Monkey patch the Marimo DataFrame formatter
+        """
+        format_mapping = {}
+
+        # Check for MoleculeDtype / DisplayDtype (OEPandas specific)
+        if oepandas_available:
+            for col in df.columns:
+                dtype = df[col].dtype
+
+                if isinstance(dtype, oepd.MoleculeDtype):
+                    arr = df[col].array
+                    ctx = get_series_context(arr.metadata).copy()
+                    format_mapping[col] = _create_molecule_formatter(ctx)
+
+                elif isinstance(dtype, oepd.DisplayDtype):
+                    arr = df[col].array
+                    ctx = get_series_context(arr.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_()
+
+    # 2. Inject into the Registry
+    def install_marimo_pandas_formatter():
+        # Check if we've already installed it to avoid duplicates
+        for typ, func in OPINIONATED_FORMATTERS.formatters.items():
+            if typ is pd.DataFrame and func.__name__ == "marimo_pandas_formatter":
+                return  # Already installed
+
+        OPINIONATED_FORMATTERS.formatters[pd.DataFrame] = marimo_pandas_formatter
+
+    # 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}')
+
+
+########################################################################################################################
+# Fallback DataFrame MIME handler for non-Marimo contexts
+########################################################################################################################
+
+def _display_dataframe(self: pd.DataFrame):
+    """
+    Fallback MIME hook for Pandas 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_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