Adding more type annotations to the codebase

There is a long discussion in gh-17945 about adding type annotations to the sympy codebase but that is more of a historical issue with many tangents. Some parts of the codebase now have type annotations and this issue is about incrementally adding more type annotations. PRs are welcome both big and small (better to start small) that add type annotations to parts of the codebase where it makes sense to do so (not necessarily anywhere in the codebase - see below).

You can see related issues and PRs under the typing label: https://github.com/sympy/sympy/issues?q=label%3A%22typing%22

PEP 484 introduced type annotations to Python but the sympy codebase predates that and so did not have type annotations when most parts were written. Over time the role of type annotations in Python has become a bit clearer and the situation is that type annotations are basically not optional for a library like SymPy especially on public interfaces. Many users of SymPy will not necessarily use type annotations much in their own code but can still benefit from type checker based tools for things like autocomplete and for checking types in their own code. You can see an examples of this in my editor:

Those are coming from the pyright type checker LSP server which is installed in my editor and can read the type annotations for Symbol, Expr.__add__ and so on. There is also pylance which is based on pyright and is part of the standard Python setup in vscode. If working on adding type annotations to the codebase I recommend having one of these installed in your editor so that you can see the type checker running interactively. In those demonstrations I chose examples that work but it is not hard to find examples that do not work because of missing type annotations in the codebase.

You can also run the mypy type checker on the file shown above:

$ mypy t.py
t.py:7: note: Revealed type is "Any"
Success: no issues found in 1 source file

Here mypy has revealed the type to be Any even though pyright understands that it is Expr. I'm not sure why exactly that happens in this case but it is not uncommon that pyright can understand what the type is but mypy cannot. Ideally we want both to understand what the type is but that is not always possible.

In CI type annotations are checked using the mypy type checker that you can run from the git repo with

$ mypy sympy

Note that mypy does not check the body of functions if the function does not have a return type annotation. What that means is that mypy does not check most of the functions in the codebase but when adding type annotations you then need to resolve the type check errors that mypy throws up which in turn likely mean adding more type annotations to other functions that are called from that function and so on. It is always necessary that any PR passes under mypy sympy because that is the CI check. Note that pyright and mypy can have different opinions about the types sometimes and that mypy fails to understand some things that are quite important in sympy.

We can use pyright to measure how much of the codebase is covered by a type checker although the package first has to be installed:

$ pip install .
$ pyright --verifytypes sympy
...
<a million error messages ...>
...

Symbols exported by "sympy": 39643
  With known type: 3474
  With ambiguous type: 1764
  With unknown type: 34405

Other symbols referenced but not exported by "sympy": 10245
  With known type: 5599
  With ambiguous type: 236
  With unknown type: 4410

Symbols without documentation:
  Functions without docstring: 17200
  Functions without default param: 104
  Classes without docstring: 1063

Type completeness score: 8.8%

Completed in 160.466sec

So pyright says that the type completeness score is 8.8% which is quite low but this is a bit of an underestimate since it treats all of the test code as being public functions (maybe there is a way to configure pyright to ignore those). I don't think it makes sense to add type annotations to test code usually but it is useful to look at test code with the type checker running in your editor since it shows how it would look for a user writing that code.

An ongoing project that will take a long time is to add more type annotations to the codebase. The goal is not necessarily to get to a 100% score from pyright but rather:

• Ensure that the public interface that users are using is sufficiently covered by type annotations so that examples like I showed in my editor above work nicely for end users.
• Make some particular parts of the codebase fully covered by type annotations and well typed in all the internal code. This applies to the polys module (gh-27360), ntheory, core, logic, functions, so that the core modules in the codebase are fully checked by type checkers.

The main way to evaluate progress with this is not the score from pyright but just whether or not type checkers are able to infer the types correctly in common code in the codebase or that users might write.

You can also try something like

$ pyright sympy/algebras/quaternion.py 
WARNING: there is a new pyright version available (v1.1.403 -> v1.1.407).
Please install the new version or set PYRIGHT_PYTHON_FORCE_VERSION to `latest`

/.../sympy/sympy/algebras/quaternion.py
  /.../sympy/sympy/algebras/quaternion.py:462:20 - error: Type "Any | Unknown | Expr | Basic | atan2" is not assignable to return type "Quaternion"
    Type "Any | Unknown | Expr | Basic | atan2" is not assignable to type "Quaternion"
      "Basic" is not assignable to "Quaternion" (reportReturnType)
  /.../sympy/sympy/algebras/quaternion.py:464:20 - error: Type "Any | Unknown | Expr | Basic | atan2" is not assignable to return type "Quaternion"
    Type "Any | Unknown | Expr | Basic | atan2" is not assignable to type "Quaternion"
      "Basic" is not assignable to "Quaternion" (reportReturnType)
  /.../sympy/sympy/algebras/quaternion.py:1242:16 - error: Type "tuple[tuple[Any | Unknown | Expr | Basic | atan2, Any | Unknown | Expr | Basic | atan2, Any | Unknown | Expr | Basic | atan2], Any | Unknown | Expr | Basic | atan2]" is not assignable to return type "tuple[tuple[Expr, Expr, Expr], Expr]"
    "tuple[tuple[Any | Unknown | Expr | Basic | atan2, Any | Unknown | Expr | Basic | atan2, Any | Unknown | Expr | Basic | atan2], Any | Unknown | Expr | Basic | atan2]" is not assignable to "tuple[tuple[Expr, Expr, Expr], Expr]"
      Tuple entry 1 is incorrect type
        "tuple[Any | Unknown | Expr | Basic | atan2, Any | Unknown | Expr | Basic | atan2, Any | Unknown | Expr | Basic | atan2]" is not assignable to "tuple[Expr, Expr, Expr]"
          Tuple entry 1 is incorrect type
            Type "Any | Unknown | Expr | Basic | atan2" is not assignable to type "Expr" (reportReturnType)
3 errors, 0 warnings, 0 informations 

That particular module has a lot of type annotations relative to other parts of the codebase so it has much fewer errors than you would see in other cases. Still though there are some errors and I think those are all because of the top-level trigsimp function not having return type annotations. Without the return annotation mypy treats the return type as Any and so does not give any errors but pyright tries to infer the return type by looking into the code and often comes up with strange answers like Any | Unknown | Expr | Basic | atan2 and then ends up reporting errors when using that type. Note that the particular example of trigsimp is a case where the types can be complicated. In most cases the input and output would be Expr but it is also valid to pass in other types like int or say a Matrix (which is not an Expr):

In [1]: M = Matrix([[1, 2], [3, 4]])

In [2]: trigsimp(M)
Out[2]: 
⎡1  2⎤
⎢    ⎥
⎣3  4⎦

In [3]: isinstance(M, Expr)
Out[3]: False

In [4]: isinstance(M, trigsimp(Expr))
Out[4]: False

That means that writing annotations for that function is complicated because it is "overloaded" so I would not recommend adding annotations to trigsimp until after gaining a lot of familiarity with using the type checkers and adding annotations in the sympy codebase specifically. Just figuring out what the actual input and output types of the trigsimp function are is nontrivial and it is not useful to just open a PR with a guess about what the types are. Instead it is better to focus at first on adding annotations to functions where the types are simple to express but are not currently inferred correctly by pyright. Also in some cases the existing type annotations may be incorrect or incomplete and fixing incorrect type annotations is especially important.

Many functions in the codebase have very strange types that are hard or perhaps even impossible to express using type annotations so in some cases it is better to add new functions that have simpler types and then change the codebase to use those instead. In many cases adding type annotations or just having a type checker running in your editor will reveal obvious bugs in the code that should then be fixed.

To be clear it might sound like adding annotations is a good job for LLMs but this is not something that is suitable for getting an LLM to do. An LLM would just guess the types which might be reasonable if the codebase was near to 100% covered by type annotations so that we could depend on the type checkers to check that the types are in fact correct. However the codebase is nowhere near 100% covered and the fact that a type checker passes does not remotely prove that the annotations are correct (many are already incorrect). If you want to add/fix the type annotations then you really need to read the code carefully with a type checker running interactively in your editor and check through the types of all variables to see if the type checker matches what the types actually are at runtime.

One thing that can be useful if you are unsure what the types are is to add something in the code like

assert isinstance(some_variable, int)

Then you can run the test suite with

pytest sympy -n auto

and see if the assert causes any of the tests to fail. Then if they do fail you can run

pytest sympy --lf --pdb

That will enter the pdb debugger at the point where the assert failed and use e.g.

(pdb) p type(some_variable)

to see what the type is.

If you are interested in working on this then go ahead but please do not just add annotations randomly to the codebase and please do not submit any LLM generated PRs. I listed a few modules above where there should probably be complete type annotations but otherwise there are some large parts of the codebase where adding type annotations is probably not welcome so ask in this issue and we can discuss where it is reasonable to add them or what the annotations for any particular function should be.

Any PR or commit for this issue should include a link to this issue but should not say that it "fixes" the issue because no single PR or commit is going to fix this and if you write "Fixes gh-1234" then when a PR is merged it closes the issue.

EDIT: Since the same mistakes are happening over and over in many different PRs let's just list a few points here:

• You need to add from __future__ import annotations at the top of any file when adding type annotations to it.
• Using e.g. typing.List[int] is the older syntax that is not needed in the sympy codebase any more. Instead it can be list[int] using the actual list type.