diff --git a/rdkit-stubs/Chem/rdchem.pyi b/rdkit-stubs/Chem/rdchem.pyi index 175dc50..7594cf2 100644 --- a/rdkit-stubs/Chem/rdchem.pyi +++ b/rdkit-stubs/Chem/rdchem.pyi @@ -1,7 +1,7 @@ from collections.abc import Iterable -from typing import Any, ClassVar, Optional, Sequence +from typing import Any, ClassVar, Literal, Optional, Sequence, overload -from typing import overload +import numpy as np ALLOW_CHARGE_SEPARATION: ResonanceFlags ALLOW_INCOMPLETE_OCTETS: ResonanceFlags @@ -425,8 +425,9 @@ class Conformer: def GetNumAtoms(cls, RDKit) -> Any: ... @classmethod def GetOwningMol(cls, RDKit) -> Any: ... - @classmethod - def GetPositions(cls, RDKit) -> Any: ... + def GetPositions( + self, + ) -> np.ndarray[tuple[int, Literal[3]], np.dtype[np.float64]]: ... @classmethod def GetProp(cls, *args, **kwargs) -> Any: ... @classmethod @@ -543,8 +544,7 @@ class Mol: def GetBonds(cls, boost) -> Any: ... @classmethod def GetBoolProp(cls, *args, **kwargs) -> Any: ... - @classmethod - def GetConformer(cls, RDKit) -> Any: ... + def GetConformer(self, id: int = -1) -> Conformer: ... @classmethod def GetConformers(cls, boost) -> Any: ... @classmethod diff --git a/rdkit-stubs/Chem/rdmolops.pyi b/rdkit-stubs/Chem/rdmolops.pyi index 81d84fb..62198f9 100644 --- a/rdkit-stubs/Chem/rdmolops.pyi +++ b/rdkit-stubs/Chem/rdmolops.pyi @@ -1,6 +1,7 @@ from typing import Any, ClassVar, Sequence from rdkit.Chem.rdchem import Mol +from rdkit.Chem.rdMolDescriptors import AtomPairsParameters from rdkit.DataStructs.cDataStructs import ExplicitBitVect ADJUST_IGNOREALL: AdjustQueryWhichFlags @@ -175,7 +176,13 @@ class _vectN5RDKit9Chirality10StereoInfoE: @classmethod def __setitem__(cls, index, object) -> Any: ... -def AddHs(RDKit) -> Any: ... +def AddHs( + mol: Mol, + explicitOnly: bool = False, + addCoords: bool = False, + onlyOnAtoms: AtomPairsParameters | None = None, + addResidueInfo: bool = False, +) -> Mol: ... def AddRecursiveQuery(*args, **kwargs) -> Any: ... def AddWavyBondsForStereoAny(RDKit) -> Any: ... def AdjustQueryProperties(RDKit) -> Any: ... diff --git a/setup.py b/setup.py index 7877b87..db3dfce 100644 --- a/setup.py +++ b/setup.py @@ -29,7 +29,7 @@ def list_package_files() -> List[str]: setup( name="rdkit-stubs", - version="0.5", + version="0.6", description="type stubs for rdkit", author="Andrew Dirksen, Ryan Rightmer", author_email="andrew@dirksen.com, rrightmer@gmail.com", diff --git a/test/test_stubs.py b/test/test_stubs.py index d0a6bf1..cc7d39a 100644 --- a/test/test_stubs.py +++ b/test/test_stubs.py @@ -41,6 +41,20 @@ def sdwriter(): writer.SetProps(("asdf",)) +def conformers(): + mol = sample_mol() + rdDistGeom.EmbedMolecule(mol) + conf = mol.GetConformer() + poses = conf.GetPositions() + assert len(poses.shape) == 2 + assert poses.shape[1] == 3 + + +def add_hs(): + mol = sample_mol() + Chem.AddHs(mol) + + def embed(): mol = sample_mol() rdDistGeom.EmbedMolecule(mol) @@ -49,4 +63,6 @@ def embed(): def test_frob(): frob(sample_mol()) sdwriter() + conformers() + add_hs() embed()