rdkit-cli 0.1.0__py3-none-any.whl

rdkit_cli/io/readers.py

@@ -0,0 +1,352 @@
+"""File readers for various molecular file formats."""
+
+import sys
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Iterator, Optional, Any
+
+import pandas as pd
+from rdkit import Chem
+
+from rdkit_cli.io.formats import FileFormat, FormatConfig, detect_format
+
+
+class MoleculeRecord:
+    """A molecule with its associated metadata."""
+
+    __slots__ = ("mol", "smiles", "name", "metadata", "row_idx")
+
+    def __init__(
+        self,
+        mol: Optional[Chem.Mol],
+        smiles: str = "",
+        name: str = "",
+        metadata: Optional[dict[str, Any]] = None,
+        row_idx: int = -1,
+    ):
+        self.mol = mol
+        self.smiles = smiles
+        self.name = name
+        self.metadata = metadata or {}
+        self.row_idx = row_idx
+
+    @property
+    def is_valid(self) -> bool:
+        """Check if molecule was parsed successfully."""
+        return self.mol is not None
+
+
+class MoleculeReader(ABC):
+    """Abstract base class for molecule file readers."""
+
+    @abstractmethod
+    def __iter__(self) -> Iterator[MoleculeRecord]:
+        """Yield MoleculeRecord objects."""
+        pass
+
+    @abstractmethod
+    def __len__(self) -> int:
+        """Return total number of molecules (for progress)."""
+        pass
+
+    @abstractmethod
+    def close(self):
+        """Close any open resources."""
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+
+class CSVReader(MoleculeReader):
+    """Read molecules from CSV/TSV files."""
+
+    def __init__(
+        self,
+        path: Path | str,
+        smiles_column: str = "smiles",
+        name_column: Optional[str] = None,
+        delimiter: str = ",",
+        has_header: bool = True,
+    ):
+        self.path = Path(path)
+        self.smiles_column = smiles_column
+        self.name_column = name_column
+        self.delimiter = delimiter
+        self.has_header = has_header
+        self._count: Optional[int] = None
+        self._df: Optional[pd.DataFrame] = None
+
+    def __len__(self) -> int:
+        if self._count is None:
+            # Count lines efficiently
+            with open(self.path, "rb") as f:
+                self._count = sum(1 for _ in f) - (1 if self.has_header else 0)
+        return self._count
+
+    def __iter__(self) -> Iterator[MoleculeRecord]:
+        header = 0 if self.has_header else None
+
+        # Read in chunks for memory efficiency
+        for chunk in pd.read_csv(
+            self.path,
+            delimiter=self.delimiter,
+            header=header,
+            chunksize=10000,
+            dtype=str,
+            na_filter=False,
+        ):
+            # Handle no-header case
+            if not self.has_header:
+                # Assume first column is SMILES
+                chunk.columns = [self.smiles_column] + [f"col_{i}" for i in range(1, len(chunk.columns))]
+
+            for idx, row in chunk.iterrows():
+                smiles = str(row.get(self.smiles_column, ""))
+                name = str(row.get(self.name_column, "")) if self.name_column else ""
+
+                mol = None
+                if smiles:
+                    try:
+                        mol = Chem.MolFromSmiles(smiles)
+                    except Exception:
+                        pass
+
+                if mol is None and smiles:
+                    print(f"Warning: Failed to parse SMILES at row {idx}: {smiles[:50]}", file=sys.stderr)
+
+                yield MoleculeRecord(
+                    mol=mol,
+                    smiles=smiles,
+                    name=name,
+                    metadata=row.to_dict(),
+                    row_idx=int(idx),
+                )
+
+    def close(self):
+        pass
+
+
+class SMIReader(MoleculeReader):
+    """Read molecules from SMILES files."""
+
+    def __init__(
+        self,
+        path: Path | str,
+        has_header: bool = False,
+        delimiter: str = " ",
+    ):
+        self.path = Path(path)
+        self.has_header = has_header
+        self.delimiter = delimiter
+        self._count: Optional[int] = None
+
+    def __len__(self) -> int:
+        if self._count is None:
+            with open(self.path, "r") as f:
+                count = sum(1 for _ in f)
+            self._count = count - (1 if self.has_header else 0)
+        return self._count
+
+    def __iter__(self) -> Iterator[MoleculeRecord]:
+        with open(self.path, "r") as f:
+            if self.has_header:
+                next(f)
+
+            for idx, line in enumerate(f):
+                line = line.strip()
+                if not line:
+                    continue
+
+                parts = line.split(self.delimiter, 1)
+                smiles = parts[0] if parts else ""
+                name = parts[1].strip() if len(parts) > 1 else ""
+
+                mol = None
+                if smiles:
+                    try:
+                        mol = Chem.MolFromSmiles(smiles)
+                    except Exception:
+                        pass
+
+                if mol is None and smiles:
+                    print(f"Warning: Failed to parse SMILES at line {idx + 1}: {smiles[:50]}", file=sys.stderr)
+
+                yield MoleculeRecord(
+                    mol=mol,
+                    smiles=smiles,
+                    name=name,
+                    metadata={"smiles": smiles, "name": name},
+                    row_idx=idx,
+                )
+
+    def close(self):
+        pass
+
+
+class SDFReader(MoleculeReader):
+    """Read molecules from SDF files."""
+
+    def __init__(self, path: Path | str):
+        self.path = Path(path)
+        self._count: Optional[int] = None
+
+    def __len__(self) -> int:
+        if self._count is None:
+            # Count $$$$ delimiters
+            with open(self.path, "rb") as f:
+                self._count = sum(1 for line in f if line.strip() == b"$$$$")
+        return self._count
+
+    def __iter__(self) -> Iterator[MoleculeRecord]:
+        supplier = Chem.SDMolSupplier(str(self.path))
+
+        for idx, mol in enumerate(supplier):
+            metadata = {}
+            smiles = ""
+            name = ""
+
+            if mol is not None:
+                # Extract properties
+                for prop in mol.GetPropsAsDict():
+                    metadata[prop] = mol.GetProp(prop)
+
+                name = mol.GetProp("_Name") if mol.HasProp("_Name") else ""
+                smiles = Chem.MolToSmiles(mol)
+            else:
+                print(f"Warning: Failed to parse molecule at index {idx}", file=sys.stderr)
+
+            yield MoleculeRecord(
+                mol=mol,
+                smiles=smiles,
+                name=name,
+                metadata=metadata,
+                row_idx=idx,
+            )
+
+    def close(self):
+        pass
+
+
+class ParquetReader(MoleculeReader):
+    """Read molecules from Parquet files."""
+
+    def __init__(
+        self,
+        path: Path | str,
+        smiles_column: str = "smiles",
+        name_column: Optional[str] = None,
+    ):
+        self.path = Path(path)
+        self.smiles_column = smiles_column
+        self.name_column = name_column
+        self._count: Optional[int] = None
+
+    def __len__(self) -> int:
+        if self._count is None:
+            import pyarrow.parquet as pq
+            self._count = pq.read_metadata(self.path).num_rows
+        return self._count
+
+    def __iter__(self) -> Iterator[MoleculeRecord]:
+        import pyarrow.parquet as pq
+
+        # Read in batches for memory efficiency
+        parquet_file = pq.ParquetFile(self.path)
+
+        row_idx = 0
+        for batch in parquet_file.iter_batches(batch_size=10000):
+            df = batch.to_pandas()
+
+            for _, row in df.iterrows():
+                smiles = str(row.get(self.smiles_column, ""))
+                name = str(row.get(self.name_column, "")) if self.name_column else ""
+
+                mol = None
+                if smiles:
+                    try:
+                        mol = Chem.MolFromSmiles(smiles)
+                    except Exception:
+                        pass
+
+                if mol is None and smiles:
+                    print(f"Warning: Failed to parse SMILES at row {row_idx}: {smiles[:50]}", file=sys.stderr)
+
+                yield MoleculeRecord(
+                    mol=mol,
+                    smiles=smiles,
+                    name=name,
+                    metadata=row.to_dict(),
+                    row_idx=row_idx,
+                )
+                row_idx += 1
+
+    def close(self):
+        pass
+
+
+def create_reader(
+    path: str | Path,
+    format_config: Optional[FormatConfig] = None,
+    smiles_column: str = "smiles",
+    name_column: Optional[str] = None,
+    has_header: Optional[bool] = None,
+) -> MoleculeReader:
+    """
+    Factory function to create appropriate reader.
+
+    Args:
+        path: Path to input file
+        format_config: Optional format configuration
+        smiles_column: Name of SMILES column (for CSV/Parquet)
+        name_column: Name of name column
+        has_header: Override header detection
+
+    Returns:
+        Appropriate MoleculeReader instance
+    """
+    path = Path(path)
+
+    if format_config is None:
+        file_format = detect_format(path)
+    else:
+        file_format = format_config.format
+        smiles_column = format_config.smiles_column
+        name_column = format_config.name_column
+        if has_header is None:
+            has_header = format_config.has_header
+
+    if file_format == FileFormat.CSV:
+        return CSVReader(
+            path,
+            smiles_column=smiles_column,
+            name_column=name_column,
+            delimiter=",",
+            has_header=has_header if has_header is not None else True,
+        )
+    elif file_format == FileFormat.TSV:
+        return CSVReader(
+            path,
+            smiles_column=smiles_column,
+            name_column=name_column,
+            delimiter="\t",
+            has_header=has_header if has_header is not None else True,
+        )
+    elif file_format == FileFormat.SMI:
+        return SMIReader(
+            path,
+            has_header=has_header if has_header is not None else False,
+        )
+    elif file_format == FileFormat.SDF:
+        return SDFReader(path)
+    elif file_format == FileFormat.PARQUET:
+        return ParquetReader(
+            path,
+            smiles_column=smiles_column,
+            name_column=name_column,
+        )
+    else:
+        raise ValueError(f"Unsupported format: {file_format}")

rdkit_cli/io/writers.py

@@ -0,0 +1,275 @@
+"""File writers for various molecular file formats."""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Optional
+
+import pandas as pd
+from rdkit import Chem
+
+from rdkit_cli.io.formats import FileFormat, detect_format
+
+
+class MoleculeWriter(ABC):
+    """Abstract base class for molecule file writers."""
+
+    @abstractmethod
+    def write_row(self, data: dict[str, Any]):
+        """Write a single row of data."""
+        pass
+
+    @abstractmethod
+    def write_batch(self, data: list[dict[str, Any]]):
+        """Write a batch of results."""
+        pass
+
+    @abstractmethod
+    def close(self):
+        """Finalize and close the file."""
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()
+
+
+class CSVWriter(MoleculeWriter):
+    """Write results to CSV/TSV files."""
+
+    def __init__(
+        self,
+        path: Path | str,
+        delimiter: str = ",",
+        columns: Optional[list[str]] = None,
+    ):
+        self.path = Path(path)
+        self.delimiter = delimiter
+        self.columns = columns
+        self._file = open(path, "w", newline="", encoding="utf-8")
+        self._header_written = False
+        self._column_order: Optional[list[str]] = None
+
+    def write_row(self, data: dict[str, Any]):
+        """Write a single row."""
+        self.write_batch([data])
+
+    def write_batch(self, data: list[dict[str, Any]]):
+        """Write a batch of results."""
+        if not data:
+            return
+
+        # Determine column order from first row if not set
+        if self._column_order is None:
+            if self.columns:
+                self._column_order = self.columns
+            else:
+                self._column_order = list(data[0].keys())
+
+        # Write header if not done yet
+        if not self._header_written:
+            self._file.write(self.delimiter.join(self._column_order) + "\n")
+            self._header_written = True
+
+        # Write data rows
+        for row in data:
+            values = []
+            for col in self._column_order:
+                val = row.get(col, "")
+                # Handle special types
+                if val is None:
+                    val = ""
+                elif isinstance(val, float):
+                    if pd.isna(val):
+                        val = ""
+                    else:
+                        val = str(val)
+                else:
+                    val = str(val)
+                # Escape delimiter and quotes
+                if self.delimiter in val or '"' in val or "\n" in val:
+                    val = '"' + val.replace('"', '""') + '"'
+                values.append(val)
+            self._file.write(self.delimiter.join(values) + "\n")
+
+    def close(self):
+        """Close the file."""
+        if self._file:
+            self._file.close()
+            self._file = None
+
+
+class SMIWriter(MoleculeWriter):
+    """Write molecules to SMILES files."""
+
+    def __init__(
+        self,
+        path: Path | str,
+        smiles_column: str = "smiles",
+        name_column: Optional[str] = "name",
+    ):
+        self.path = Path(path)
+        self.smiles_column = smiles_column
+        self.name_column = name_column
+        self._file = open(path, "w", encoding="utf-8")
+
+    def write_row(self, data: dict[str, Any]):
+        """Write a single row."""
+        smiles = data.get(self.smiles_column, "")
+        name = data.get(self.name_column, "") if self.name_column else ""
+
+        if smiles:
+            if name:
+                self._file.write(f"{smiles} {name}\n")
+            else:
+                self._file.write(f"{smiles}\n")
+
+    def write_batch(self, data: list[dict[str, Any]]):
+        """Write a batch of results."""
+        for row in data:
+            self.write_row(row)
+
+    def close(self):
+        """Close the file."""
+        if self._file:
+            self._file.close()
+            self._file = None
+
+
+class SDFWriter(MoleculeWriter):
+    """Write molecules to SDF files."""
+
+    def __init__(self, path: Path | str):
+        self.path = Path(path)
+        self._writer = Chem.SDWriter(str(path))
+
+    def write_row(self, data: dict[str, Any]):
+        """Write a single row."""
+        mol = data.get("mol")
+
+        if mol is None:
+            # Try to create from SMILES
+            smiles = data.get("smiles", "")
+            if smiles:
+                mol = Chem.MolFromSmiles(smiles)
+
+        if mol is not None:
+            # Set properties from data
+            for key, value in data.items():
+                if key not in ("mol", "smiles") and value is not None:
+                    try:
+                        mol.SetProp(str(key), str(value))
+                    except Exception:
+                        pass
+            self._writer.write(mol)
+
+    def write_batch(self, data: list[dict[str, Any]]):
+        """Write a batch of results."""
+        for row in data:
+            self.write_row(row)
+
+    def close(self):
+        """Close the writer."""
+        if self._writer:
+            self._writer.close()
+            self._writer = None
+
+
+class ParquetWriter(MoleculeWriter):
+    """Write results to Parquet files."""
+
+    def __init__(
+        self,
+        path: Path | str,
+        columns: Optional[list[str]] = None,
+    ):
+        self.path = Path(path)
+        self.columns = columns
+        self._batches: list[dict[str, Any]] = []
+        self._batch_size = 100000  # Write in batches of 100k
+
+    def write_row(self, data: dict[str, Any]):
+        """Write a single row."""
+        # Remove mol objects (not serializable)
+        clean_data = {k: v for k, v in data.items() if k != "mol"}
+        self._batches.append(clean_data)
+
+        if len(self._batches) >= self._batch_size:
+            self._flush()
+
+    def write_batch(self, data: list[dict[str, Any]]):
+        """Write a batch of results."""
+        for row in data:
+            clean_data = {k: v for k, v in row.items() if k != "mol"}
+            self._batches.append(clean_data)
+
+        if len(self._batches) >= self._batch_size:
+            self._flush()
+
+    def _flush(self):
+        """Write accumulated batches to file."""
+        if not self._batches:
+            return
+
+        import pyarrow as pa
+        import pyarrow.parquet as pq
+
+        df = pd.DataFrame(self._batches)
+
+        # Reorder columns if specified
+        if self.columns:
+            cols = [c for c in self.columns if c in df.columns]
+            extra = [c for c in df.columns if c not in self.columns]
+            df = df[cols + extra]
+
+        table = pa.Table.from_pandas(df, preserve_index=False)
+
+        if self.path.exists():
+            # Append to existing file
+            existing = pq.read_table(self.path)
+            table = pa.concat_tables([existing, table])
+
+        pq.write_table(table, self.path)
+        self._batches = []
+
+    def close(self):
+        """Finalize and close the file."""
+        self._flush()
+
+
+def create_writer(
+    path: str | Path,
+    format_override: Optional[FileFormat] = None,
+    columns: Optional[list[str]] = None,
+    smiles_column: str = "smiles",
+    name_column: Optional[str] = "name",
+) -> MoleculeWriter:
+    """
+    Factory function to create appropriate writer.
+
+    Args:
+        path: Output file path
+        format_override: Override auto-detected format
+        columns: Column order for output
+        smiles_column: Name of SMILES column (for SMI files)
+        name_column: Name of name column (for SMI files)
+
+    Returns:
+        Appropriate MoleculeWriter instance
+    """
+    path = Path(path)
+    file_format = format_override or detect_format(path)
+
+    if file_format == FileFormat.CSV:
+        return CSVWriter(path, delimiter=",", columns=columns)
+    elif file_format == FileFormat.TSV:
+        return CSVWriter(path, delimiter="\t", columns=columns)
+    elif file_format == FileFormat.SMI:
+        return SMIWriter(path, smiles_column=smiles_column, name_column=name_column)
+    elif file_format == FileFormat.SDF:
+        return SDFWriter(path)
+    elif file_format == FileFormat.PARQUET:
+        return ParquetWriter(path, columns=columns)
+    else:
+        raise ValueError(f"Unsupported format: {file_format}")

rdkit_cli/parallel/__init__.py

@@ -0,0 +1,5 @@
+"""Parallel processing utilities."""
+
+from rdkit_cli.parallel.executor import ParallelExecutor, parallel_map
+
+__all__ = ["ParallelExecutor", "parallel_map"]