Add type stubs for sympy (#307)

* Initial commit

* Update readme

* Add directions on how created

* More directions

* More directions

* Typo

README.md

@@ -54,7 +54,6 @@ Stubs for the following libraries now exist in typeshed or the libraries themsel
 - retry
 - slugify
 - SQLAlchemy (see <https://pypi.org/project/types-SQLAlchemy/> for SQLAlchemy 1.4; 2.0.0 and above include type annotations)
-- sympy
 - tenacity
 
 The following libraries are `py.typed`. We still have stubs for them here, but we are no longer actively maintaining them. We continue to bundle them with Pylance so users on older, non-`py.typed` versions will still get type info. If you find problems in our stubs for these libraries, rather than filing an issue here, you should upgrade to the version shown below to get the official stubs:

stubs/sympy/README.md

@@ -0,0 +1,156 @@
+These stubs were generated by:
+
+- Running `pyright --createstub sympy` after installing sympy into a venv
+- Running a script that uncommented all of the return values in the generated stubs
+- Running a script that removed all doc strings from the stubs (as sympy has docstrings itself)
+- Adding a `partial` py.typed file
+- Fixing all import errors one at a time (using Pylance)
+
+The last part took the longest. It might be quicker to just add any changes by hand, rather than regenerating from scratch.
+
+
+Scripts for future use:
+
+### Uncomment return values
+
+```python
+import os
+import re
+import shutil
+from typing import List
+
+SOURCE_DIR = "typings\\sympy"
+DEST_DIR = "typings\\sympy-returnvalues"
+
+def read_file(file_path: str) -> List[str]:
+    with open(file_path, "r") as file:
+        return file.readlines()
+
+def write_file(file_path: str, lines: List[str]) -> None:
+    try:
+        os.makedirs(os.path.dirname(file_path))
+    except FileExistsError:
+        pass
+    with open(file_path, "w") as file:
+       file.writelines(lines)
+
+def fix_file(file_path: str, dest_path: str, sub_package: str) -> None:
+    lines = read_file(file_path)[4:]
+
+    # Turn the return type comments into actual return types
+    changed = [re.sub(r":\s+#(.*)", r"\1", line) for line in lines]
+
+    # Replace `from .` imports with `from sympy.<subpackage>.`
+    replace_str = f"from sympy.{sub_package + '.' if sub_package else ''}"
+    changed = [re.sub(r"from \.", replace_str, line) for line in changed]
+    any_changes = [line for line in changed if line not in lines]
+
+    if len(any_changes) > 0:
+        # Find this same file in the '.venv/lib/site-packages/sympy' directory and see if it has typing information in it.
+        site_packages_file = file_path.replace(SOURCE_DIR, ".venv\\lib\\site-packages\\sympy").replace(".pyi", ".py")
+        if os.path.exists(site_packages_file):
+            site_packages_lines = read_file(site_packages_file)
+            if "->" in site_packages_lines:
+                return
+            else:
+                print(f"Writing {dest_path}")
+                write_file(dest_path, changed)
+
+def fix_all_stubs() -> None:
+    stubs_dir = SOURCE_DIR
+    dest_dir = DEST_DIR
+    shutil.rmtree(dest_dir, ignore_errors=True)
+    os.makedirs(dest_dir, exist_ok=True)
+    # First write a partial py.typed file to the destination directory
+    write_file(os.path.join(dest_dir, "py.typed"), ["partial"])
+
+    # Then iterate over all of the generated files and fix them up so they're valid
+    for root, dirs, files in os.walk(stubs_dir):
+        for file in files:
+            if file.endswith(".pyi"):
+                file_path = os.path.join(root, file)
+                dest_path = file_path.replace(stubs_dir, dest_dir)
+                sub_dir_pos = root.index(stubs_dir) + len(stubs_dir) + 1
+                sub_dir = root[sub_dir_pos:]
+                sub_package = sub_dir.replace("\\", ".")
+                fix_file(file_path, dest_path, sub_package)
+
+fix_all_stubs()
+
+```
+
+### Remove doc comments
+```python
+
+from codecs import ignore_errors
+import os
+import shutil
+from typing import List
+
+
+SOURCE_DIR = "typings\\sympy-returnvalues"
+DEST_DIR = "typings\\sympy-docs"
+
+
+def fix_file(file_path:str, dest_path:str):
+    # Read the file one char at a time until we find one of r""", b""", or """
+    with open(file_path, mode="r") as file:
+        lines = file.readlines()
+    contents = "".join(lines)
+    new_contents = ""
+    in_docstring = False
+    ignoring_line = False
+    line_start = 0
+    i = 0
+    while i < len(contents):
+        char = contents[i]
+        if contents[i] == "\n":
+            line_start = len(new_contents)
+            if ignoring_line and not in_docstring:
+                ignoring_line = False
+
+        if contents[i:i+3] == "\"\"\"" or contents[i:i+3] == "\'\'\'":
+            in_docstring = not in_docstring
+            new_contents = new_contents[:line_start]
+            ignoring_line = True
+            i += 3
+        elif contents[i:i+4] == "r\"\"\"" or contents[i:i+4] == "r\'\'\'" and not in_docstring:
+            in_docstring = True
+            new_contents = new_contents[:line_start]
+            ignoring_line = True
+            i += 4
+        elif not in_docstring and not ignoring_line:
+            new_contents += char
+            i += 1
+        else:
+            i += 1
+    try:
+        os.makedirs(os.path.dirname(dest_path))
+    except FileExistsError:
+        pass
+
+    print(f"Writing {dest_path}")
+    with open(dest_path, mode="w") as file:
+        file.write(new_contents)
+
+
+def fix_all_stubs() -> None:
+    stubs_dir = SOURCE_DIR
+    dest_dir = DEST_DIR
+    shutil.rmtree(dest_dir, ignore_errors=True)
+    os.makedirs(dest_dir, exist_ok=True)
+
+    # Then iterate over all of the generated files and fix them up so they're valid
+    for root, dirs, files in os.walk(stubs_dir):
+        for file in files:
+            if file.endswith(".pyi"):
+                file_path = os.path.join(root, file)
+                dest_path = file_path.replace(stubs_dir, dest_dir)
+                sub_dir_pos = root.index(stubs_dir) + len(stubs_dir) + 1
+                sub_dir = root[sub_dir_pos:]
+                sub_package = sub_dir.replace("\\", ".")
+                fix_file(file_path, dest_path)
+
+fix_all_stubs()
+```
+

stubs/sympy/algebras/__init__.pyi

@@ -0,0 +1,3 @@
+from sympy.algebras.quaternion import Quaternion
+
+__all__ = ["Quaternion"]

stubs/sympy/algebras/quaternion.pyi

@@ -0,0 +1,173 @@
+from types import NotImplementedType
+from typing import Any, Literal, Self
+from sympy.core.basic import Basic
+from sympy.core.expr import Expr
+from sympy.core.function import UndefinedFunction
+from sympy.core.power import Pow
+from sympy.matrices.dense import MutableDenseMatrix
+
+class Quaternion(Expr):
+    _op_priority = ...
+    is_commutative = ...
+    def __new__(cls, a=..., b=..., c=..., d=..., real_field=..., norm=...) -> Self:
+        ...
+
+    def set_norm(self, norm) -> None:
+        ...
+
+    @property
+    def a(self):
+        ...
+
+    @property
+    def b(self):
+        ...
+
+    @property
+    def c(self):
+        ...
+
+    @property
+    def d(self):
+        ...
+
+    @property
+    def real_field(self):
+        ...
+
+    @property
+    def product_matrix_left(self) -> MutableDenseMatrix:
+        ...
+
+    @property
+    def product_matrix_right(self) -> MutableDenseMatrix:
+        ...
+
+    def to_Matrix(self, vector_only=...) -> MutableDenseMatrix:
+        ...
+
+    @classmethod
+    def from_Matrix(cls, elements) -> Quaternion:
+        ...
+
+    @classmethod
+    def from_euler(cls, angles, seq) -> Any:
+        ...
+
+    def to_euler(self, seq, angle_addition=..., avoid_square_root=...) -> tuple[int, ...]:
+        ...
+
+    @classmethod
+    def from_axis_angle(cls, vector, angle) -> Self:
+        ...
+
+    @classmethod
+    def from_rotation_matrix(cls, M) -> Quaternion:
+        ...
+
+    def __add__(self, other) -> Quaternion:
+        ...
+
+    def __radd__(self, other) -> Quaternion:
+        ...
+
+    def __sub__(self, other) -> Quaternion:
+        ...
+
+    def __mul__(self, other) -> Quaternion:
+        ...
+
+    def __rmul__(self, other) -> Quaternion:
+        ...
+
+    def __pow__(self, p) -> NotImplementedType | Quaternion | Literal[1]:
+        ...
+
+    def __neg__(self) -> Quaternion:
+        ...
+
+    def __truediv__(self, other):
+        ...
+
+    def __rtruediv__(self, other):
+        ...
+
+    def diff(self, *symbols, **kwargs) -> Self:
+        ...
+
+    def add(self, other) -> Quaternion:
+        ...
+
+    def mul(self, other) -> Quaternion:
+        ...
+
+    def norm(self) -> Pow:
+        ...
+
+    def normalize(self):
+        ...
+
+    def inverse(self):
+        ...
+
+    def pow(self, p) -> NotImplementedType | Quaternion | Literal[1]:
+        ...
+
+    def exp(self) -> Quaternion:
+        ...
+
+    def pow_cos_sin(self, p):
+        ...
+
+    def integrate(self, *args) -> Quaternion:
+        ...
+
+    @staticmethod
+    def rotate_point(pin, r) -> tuple[Any, Any, Any]:
+        ...
+
+    def to_axis_angle(self) -> tuple[tuple[Any, Any, Any], Any]:
+        ...
+
+    def to_rotation_matrix(self, v=..., homogeneous=...) -> MutableDenseMatrix:
+        ...
+
+    def scalar_part(self):
+        ...
+
+    def vector_part(self) -> Quaternion:
+        ...
+
+    def axis(self) -> Quaternion:
+        ...
+
+    def is_pure(self):
+        ...
+
+    def is_zero_quaternion(self) -> bool | None:
+        ...
+
+    def angle(self) -> type[UndefinedFunction]:
+        ...
+
+    def arc_coplanar(self, other) -> bool | None:
+        ...
+
+    @classmethod
+    def vector_coplanar(cls, q1, q2, q3) -> Any | bool | None:
+        ...
+
+    def parallel(self, other):
+        ...
+
+    def orthogonal(self, other):
+        ...
+
+    def index_vector(self) -> Quaternion:
+        ...
+
+    def mensor(self) -> type[UndefinedFunction]:
+        ...
+
+
+

stubs/sympy/assumptions/__init__.pyi

@@ -0,0 +1,6 @@
+from sympy.assumptions.assume import AppliedPredicate, AssumptionsContext, Predicate, assuming, global_assumptions
+from sympy.assumptions.ask import Q, ask, register_handler, remove_handler
+from sympy.assumptions.refine import refine
+from sympy.assumptions.relation import AppliedBinaryRelation, BinaryRelation
+
+__all__ = ['AppliedPredicate', 'Predicate', 'AssumptionsContext', 'assuming', 'global_assumptions', 'Q', 'ask', 'register_handler', 'remove_handler', 'refine', 'BinaryRelation', 'AppliedBinaryRelation']