Adds typing to expression tree

[pipliggins]

Description

Adds type hints to the majority of files within the expression tree.

Fixes #3497

Type of change

Please add a line in the relevant section of CHANGELOG.md to document the change (include PR #) - note reverse order of PR #s. If necessary, also add to the list of breaking changes.

• New feature (non-breaking change which adds functionality)

Key checklist:

• No style issues: $ pre-commit run (or $ nox -s pre-commit) (see CONTRIBUTING.md for how to set this up to run automatically when committing locally, in just two lines of code)
• All tests pass: $ python run-tests.py --all (or $ nox -s tests)
• The documentation builds: $ python run-tests.py --doctest (or $ nox -s doctests)

[Saransh-cpp]

@pipliggins you might want to merge develop. We introduced a few changes to the configuration of pre-commit and ruff recently.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (d201899) 99.60% compared to head (38a3713) 99.59%.
Report is 1 commits behind head on develop.

❗ Current head 38a3713 differs from pull request most recent head eeceaa7. Consider uploading reports for the commit eeceaa7 to get more accurate results

Additional details and impacted files

@@             Coverage Diff             @@
##           develop    #3578      +/-   ##
===========================================
- Coverage    99.60%   99.59%   -0.01%     
===========================================
  Files          258      258              
  Lines        21226    20865     -361     
===========================================
- Hits         21142    20781     -361     
  Misses          84       84              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[Saransh-cpp]

Thanks for working on this! Minor comments below (will go through the type hints too).

[pipliggins]

We could consider the inclusion of a MyPy check in the CI pipeline at this point.

It can be included as a pre-commit hook -

- repo: https://github.com/pre-commit/mirrors-mypy
  rev: v1.8.0
  hooks:
    - id: mypy
      files: pybamm
      args: []
      additional_dependencies:
        - numpy
        - packaging

but will this fail for the expressions that don't have a type hint at the moment? I really don't remember if it fails or if it passes (with or without the strict mode).

Anything that doesn't have a type hint mypy will ignore (more accurately, it assumes it has type "Any") without "strict" mode. If "strict" is on it will fail for the untyped expressions.

MyPy themselves acknowledge that the "strict" flag will probably be too aggressive for scenarios like this where we're adding typing to a large existing codebase. They have a number of suggestions around how to approach this here

[kratman]

Is it possible to avoid these type of ignore comments for the typing?

[pipliggins]

These comments are used by mypy when static type checking. I've done my best to use them as little as possible, but in some cases the amount of rewriting required to make mypy happy would make the code less readable than having the comment (in my opinion). If we decide not to use a type checker, they can be removed.

[kratman]

Yeah I would say this is probably a symptom of not enforcing types. What does the output say if you neglect this comment? If MyPy is not strict, is it something that we can leave as a warning to remind us to deal with it?

[pipliggins]

This instance isn't actually a brilliant example as I've managed to get rid of the type error; however there are ~ 10 ignore statements I've had to leave on a line-by-line basis. Mypy raises an error for each # type: ignore[] comment even with strict mode turned off, but if we don't enforce mypy as having to pass in the CI then we could chose to use it as a 'todo' list of sorts and remove the # type: ignore comments.

The remaining errors are split about 50/50 between difficulties integrating mypy into heavily dynamically-typed areas (e.g. it won't let you re-define a variable with a different type before you've used it, even with the --allow-redefinition flag https://mypy.readthedocs.io/en/latest/command_line.html#miscellaneous-strictness-flags), and areas where pybamm would require more significant re-writing. For example in serialise.py, where mypy complains that BaseModel doesn't have a bounds attribute. This is true at the point where BaseModel is defined, but the attribute is added once the model is discretised; as that particular function only takes a discretised model, the code works properly but I haven't found a way to flag to mypy that this is the case. There are multiple ways to fix this, including creating a DiscretisedModel class or editing the __init__ of BaseModel to include attributes that are added later, but that's a lot of work & discussion that shouldn't be done here.

[kratman]

Why do we need to enable type checking with a global variable?

[pipliggins]

Related to above - this is a statement used by mypy. In this case sympy is used in the file only as a type hint, not during runtime. TYPE_CHECKING is a constant assumed to be True by static type checkers, but is false at runtime; so sympy will only be imported if a type checker like mypy is being run. https://docs.python.org/3/library/typing.html#typing.TYPE_CHECKING

[kratman]

Personally I would prefer to just have sympy always imported. If it is used in the type annotations, then why not just always have it there?

[Saransh-cpp]

That would make sympy a required dependency. Sympy is heavy and I would avoid making it a required dependency (only a subset of features in PyBaMM require sympy)

[kratman]

I would say that if expression tree is using sympy types, then it is really a dependency. I have not downloaded this branch locally to double check, but I would guess that my IDE would complain if I did not have sympy installed too.

As we add type checking it is going to expose more things that should probably be real dependencies rather than optional ones.

[pipliggins]

sympy is only used in the expression tree for to_equation() to use with LaTeX I believe, rather than the solvers, hence the optional dependency.

[kratman]

Yeah my point was more that we might have to rethink our optional dependencies if we are using static types. Expression tree depends on sympy if it is using sympy types, but by putting this check here we are pretending that it doesn't depend on sympy.

It is too much for this task, but if we want a typed library and not have all the dependencies for types then we might need to think about how to better separate our subcomponents. Maybe things like the latexify need to go into a separate module if we want it's dependencies to be truly optional

[kratman]

We could consider the inclusion of a MyPy check in the CI pipeline at this point. However, I propose that we refrain from making it a mandatory requirement for now - we can just use it to issue warnings to users about potential issues.

Using mypy quickly escalates, we should use caution when considering using it. I tend to consider types as nice-to-haves in python, but I am certain that there are a lot of types that would need to be fixed if we use it

[pipliggins]

We could consider the inclusion of a MyPy check in the CI pipeline at this point. However, I propose that we refrain from making it a mandatory requirement for now - we can just use it to issue warnings to users about potential issues.

Using mypy quickly escalates, we should use caution when considering using it. I tend to consider types as nice-to-haves in python, but I am certain that there are a lot of types that would need to be fixed if we use it

I would tend to agree with this, but think there are ways to use it to aid consistent type hints without being bullied into making lots of changes & fixes immediately. I've linked this comment in the issue attached to the PR, I suggest discussions on the pros/cons of adding mypy are carried out there :)

[martinjrobins]

Thanks @pipliggins, great to have more types in pybamm! My main comment is that some of the very common types used are repeated in many places. Instead they should be defined in one place and then used everywhere else. e.g. you have done this in the binary ops file with ChildValue, but this should be extended to other files as well (I also think ChildValue should be ChildSymbol, as there should be another value type ChildValue used for the evaluations.

Similary you should define a single type for domain, domains, auxilary_domains as these are repeated in a lot of places. Also the types for the evaluate and _base_evaluate are repeated in a lot of places.

[martinjrobins]

great stuff, happy to merge now @pipliggins