specify type of EmbedMolecule (#5)

rdkit-stubs/Chem/rdDistGeom.pyi

@@ -1,5 +1,8 @@
 from typing import Any, ClassVar
 
+from rdkit.Chem import Mol
+from rdkit.Geometry import Point3D
+
 class EmbedParameters:
     __instance_size__: ClassVar[int] = ...
     ETversion: Any
@@ -38,7 +41,26 @@ def ETDG() -> Any: ...
 def ETKDG() -> Any: ...
 def ETKDGv2() -> Any: ...
 def ETKDGv3() -> Any: ...
-def EmbedMolecule(RDKit) -> Any: ...
+def EmbedMolecule(
+    mol: Mol,
+    maxAttempts: int = 0,
+    randomSeed: int = -1,
+    clearConfs: bool = True,
+    useRandomCoords: bool = False,
+    boxSizeMult: float = 2.0,
+    randNegEig: bool = True,
+    numZeroFail: int = 1,
+    coordMap: dict[int, Point3D] = {},
+    forceTol: float = 0.001,
+    ignoreSmoothingFailures: bool = False,
+    enforceChirality: bool = True,
+    useExpTorsionAnglePrefs: bool = True,
+    useBasicKnowledge: bool = True,
+    printExpTorsionAngles: bool = False,
+    useSmallRingTorsions: bool = False,
+    useMacrocycleTorsions: bool = False,
+    ETversion: int = 1,
+) -> int: ...
 def EmbedMultipleConfs(*args, **kwargs) -> Any: ...
 def GetMoleculeBoundsMatrix(RDKit) -> Any: ...
 def KDG() -> Any: ...

test/test_stubs.py

@@ -1,7 +1,7 @@
-from rdkit import Chem
 import io
 
-from rdkit.Chem import Mol, SDWriter
+from rdkit import Chem
+from rdkit.Chem import Mol, SDWriter, rdDistGeom
 
 
 def sample_mol() -> Mol:
@@ -41,6 +41,12 @@ def sdwriter():
     writer.SetProps(("asdf",))
 
 
+def embed():
+    mol = sample_mol()
+    rdDistGeom.EmbedMolecule(mol)
+
+
 def test_frob():
     frob(sample_mol())
     sdwriter()
+    embed()