Improve some type hints to bump mypy pin

[maresb]

Closes #6282

What is this PR about?

As explained in #6282, we were failing with the latest version of mypy for the pre-commit check, so we pinned to an older version of mypy where we were passing. This is my attempt to fix enough type issues so that we can unpin mypy and close #6282.

• Please feel free to commit directly to my branch. (I don't have much more time to put into this right now.)
• I'd be happy to unpin mypy=0.982 in a separate PR, have all my commits squashed, or for me to squash all the typing commits (leaving two commits: "improve type hints" and "unpin mypy")

I put in some medium effort to infer the correct types (but without a debugger), and probably I have a few incorrect types which should ideally be corrected. In the process of trying to infer types, I recursed in various places while adding type hints along the way. Thus this PR isn't minimal for closing #6282, and there are also files where I have updated some annotations without updating other similar annotations.

I'd like to think that things are better than they were before, and this is just about type hints which should be irrelevant at runtime, so probably this doesn't require an extremely thorough review. Nevertheless, here are several technical details to explain what I'm thinking:

• The most mysterious new errors ("notes") are:

  pymc/backends/report.py:36: note: By default the bodies of untyped functions are not checked, consider using --check-untyped-defs  [annotation-unchecked]
  

  They generally seem to point to some seemingly random line inside of a function body. Adding a return type hint to the function declaration seems to resolve the error. (Example: add -> None.)
• Type hints are much nicer after from __future__ import annotations, but I have noticed that there are currently no such imports in main. Is there any particular reason for this? (Just in case, I installed Python 3.8.0 and was successfully able to import all the relevant modules, so it seems to me like it should be fine.)
• Using collections.namedtuple is out-of-favor. It's preferable to define named tuples as classes inheriting from typing.NamedTuple, and then type annotations are provided with a nice syntax.
• Importing __future__.annotations allows for replacing Optional[SomeType] with SomeType | None and List[SomeType] with list[SomeType], and indeed pyupgrade enforces this as soon as __future__.annotations is imported.
• I don't understand why pylint was still complaining about an unused import after I added # pylint: disable=unused-import. But then mypy began complaining about the assertion assert Distribution which was meant to "use" the import. Thus I replaced the assert with a null assignment _ = Distribution to achieve the same effect without the warning.
• I gave up on these lines while trying to infer types for the named tuple State, so I sloppily set these types to Any so that I didn't have to deal with them. I'm also not very confident about q, p, and v, as I'm not so sure when what should be a RaveledVars, a numpy array, an Aesara tensor, or potentially unions of these. Any insights @aseyboldt?
• The integrator attribute for _Tree has type integration.CpuLeapfrogIntegrator, since it seems to be the only available implementation of an integrator, and there doesn't yet seem to exist an abstract Integrator class.
• For velocity I added two @overloads. Note that @overloads are type annotations and not real definitions. They are telling type checkers that when the parameter out is None, then velocity returns a numpy array, and when out is a numpy array, then velocity returns None. What doesn't seem to work is type inference from PyRight in VS Code based on the default parameter value out=None, presumably because this default may be overridden in implementing subclasses. Thus when calling velocity() I set out=None explicitly, so that VS Code infers the right types.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• ...

Bugfixes / New features

• ...

Docs / Maintenance

• Improve type hints
• Unpin mypy

[codecov]

Codecov Report

Merging #6294 (dc3abb2) into main (5d7283e) will decrease coverage by 16.77%.
The diff coverage is 97.61%.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #6294       +/-   ##
===========================================
- Coverage   94.12%   77.34%   -16.78%     
===========================================
  Files         111      111               
  Lines       23861    23901       +40     
===========================================
- Hits        22458    18486     -3972     
- Misses       1403     5415     +4012     

Impacted Files	Coverage Δ
pymc/step_methods/hmc/quadpotential.py	80.62% <80.00%> (-0.07%)	⬇️
pymc/backends/report.py	80.00% <100.00%> (ø)
pymc/blocking.py	95.34% <100.00%> (+0.22%)	⬆️
pymc/gp/util.py	96.55% <100.00%> (ø)
pymc/step_methods/hmc/base_hmc.py	90.75% <100.00%> (+0.84%)	⬆️
pymc/step_methods/hmc/hmc.py	92.98% <100.00%> (+0.25%)	⬆️
pymc/step_methods/hmc/integration.py	81.66% <100.00%> (+2.82%)	⬆️
pymc/step_methods/hmc/nuts.py	97.31% <100.00%> (+0.05%)	⬆️
pymc/variational/callbacks.py	96.22% <100.00%> (+0.22%)	⬆️
pymc/variational/operators.py	93.18% <100.00%> (+0.32%)	⬆️
... and 27 more

[maresb]

Pinging @cluhmann since this came out of yesterday's PyMC Study Group.

[cluhmann]

That looks like a tremendous effort. I know that mypy, like lots of the "housekeeping" bits of CI can be kind of annoying and that there is a constant temptation to just work around them (#6282). So thank you for putting the effort into trying to set things right (and going beyond just what was needed to prevent mypy from failing).

[aseyboldt]

This looks nice! I added a suggestion about the state in integration.py

[maresb]

Thanks so much @cluhmann for the kind words. It's nice to be appreciated. My current effort here just scratches the surface of the remaining work to be done.

Thanks a lot for the type hints @aseyboldt, that's really helpful! There seem to be a few further issues as a result, I'll see if I can fix those now...

[maresb]

Ok, I think I got it.

@aseyboldt, could you please check if 8b127a6 is safe? The np.array cast doesn't look to me like it has function. In a Python session, Python seems to convert to float as soon as any arithmetic is done, and having it as an array confuses the type checker. (In case you're worried about np.inf, I get type(np.inf) is a regular Python float, but I wouldn't be surprised if older numpies were different.)

[aseyboldt]

Good question. There are some differences between float and the numpy types:

np.array(1.) / (1. / np.inf)  # -> warning + nan
1. / (1. / np.inf)  # -> ZeroDivisionError

But I can't see anything like this happening with the energy, we only want to add, compare and store it in the stats, so I think this is safe.

[maresb]

Ooh, that's really interesting! Thanks for sharing. And thanks for confirming the safety.

Any further suggestions? Not sure what's up with codecov, but that's usually flaky, right?

[michaelosthege]

Great work!

I think we could (should?) keep mypy pinned, because we have a dependabot that will open PRs for the bumps.
This would reduce the chance of contributors suddenly running into errors with the dev environment..

Not sure what's up with codecov, but that's usually flaky, right?

Yes, it's flaky

[michaelosthege]

Shouldn't out get a real type here? Because None is not a type..

[maresb]

You're technically correct. Although the singleton None has type NoneType, I believe that type checkers automatically do an implicit conversion here. (Who really wants to remember to type the extra letters and be so pedantic? There's no potential for ambiguity since it's a singleton type.)

If you check the mypy docs, they always use None, so it's not a problem.

[maresb]

I think we could (should?) keep mypy pinned, because we have a dependabot that will open PRs for the bumps.
This would reduce the chance of contributors suddenly running into errors with the dev environment.

Sounds good. I just repinned mypy to the latest v0.990.