stubs for GetConformer, GetPositions, and AddHs (#6)

* stubs for GetConformer, GetPositions, and AddHs

* bump version

* add tests for new stubs

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

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: ...

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="[email protected], [email protected]",

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()