Configure mypy and typing stubs for custom_ops, fix errors

[mcwitt]

• Adds mypy check to pre-commit. This adds linting based on type annotations, e.g. ensuring annotations are correct and catching logic errors exposed by annotations. Initially configured with the following options:
  • ignore_missing_imports = True -- suppresses errors due to 3rd-party imports that don't have type annotations (e.g. rdkit)
  • check_untyped_defs = False -- do not run checks inside of functions that don't have type annotations (useful for code that relies on dynamic typing)
• Fixes errors exposed by the mypy check, e.g.
  • rename fep-benchmark -> fep_benchmark (the former is invalid as a module name, so it would not have worked with importlib)
  • use correct type for SimulationResult (now possible since Merges estimator_abfe.py and estimator.py #700)
  • DualTopology extends BaseTopology
  • Fixes for DeterministicMTMMove
• Adds auto-generated stub file for the custom_ops extension module and updates the CMake config to generate the stub file using mypy-stubgen
  • the stub file contains signatures extracted from the compiled .so file, and enables typechecking and auto-completion for timemachine.lib.custom_ops
  • Note that the docstrings defined in wrap_kernels.cpp are currently not included. There is an alternative to mypy-stubgen, pybind11-stubgen, that does include docstrings in the stubs -- however, there were other issues with this (e.g. needing to replace s/timemachine:://g for syntactically-valid stubs), and it seemed better to stick with the official mypy-stubgen
  • Adds pre-commit check that the stub file is up-to-date (example failure output) and a step in the CI unit-tests stage to verify this Note: the pre-commit check for stubs is skipped in CI, because it requires the .so file to be built, and this is not feasible in the linting stage. Instead, this adds a check to unit-tests, after the module has been built building the extension module as part of a pre-commit check does work and doesn't seem as unreasonable as I'd originally thought -- this is the method the PR is currently using

Review notes

• Most of additions are from the auto-generated custom_ops.pyi file. This probably doesn't need careful review, as it's only used for typechecking and completions

• Even if we choose not to adopt mypy as a pre-commit check, we might consider merging the generated stub file and the annotation fixes. The mypy check should be fully contained in 157cd7f.

Some common warnings and fixes

Using the same name for variables of different types

See docs

a = [1, 2, 3]
a = np.array(a)

Fix

a_ = [1, 2, 3]
a = np.array(a_)

Initializing empty container and appending

See docs

a = []
a.append(1)

Fix

a: List[int] = []
a.append(1)

or

a: List[Any] = []
a.append(1)

or (not recommended)

a: Any = []
a.append(1)

[maxentile]

This looks useful, independent of whether we adopt mypy

[mcwitt]

I ended up switching to use stubgen bundled with mypy; unfortunately this doesn't include docstrings defined in the pybind11 wrapping code, but seemed more reliable in my testing (for example, pybind11-stubgen generated syntactically-invalid .pyi files with timemachine:: in types)

[mcwitt]

This seems unfortunate, since we already specify the mypy dependency in 2 places (dev extras in setup.py and build system depends in pyproject.toml), but neither of these can be used here because I think we want to avoid building the package (this would add significant time to the linting stage)

[badisa]

Looks like this is no longer required, since it is in the docker container?

[mcwitt]

Good catch; removed in 829e253

[mcwitt]

Needed to add mypy here because we now use stubgen (part of mypy) in the CMake build to generate a stub file. Dependencies listed here are installed before the extension module builds, while those in setup.py are installed after.

[badisa]

If we are going to expand this dockerfile to support more CI logic, we should create a separate layer just for CI. Otherwise right now these ci requirements will get baked into production containers

[mcwitt]

This sounds good to me. We could also move over the pre-commit installation. Do you think it's worth doing in this PR or a separate one? (not sure how much work would be involved in updates to CI infra config)

[badisa]

Should be relatively straight forward, just replicate the 'final' docker stage and run the two CI specific commands within it and change the docker target in CI.

Seems best to do it in this PR, unless it isn't as trivial as I am imagining

[mcwitt]

Good call, done in 5d36eb4 (and fixed ordering in ba2d125)

[badisa]

nit: Does this need to be checked in? It seems like CMAKE will create this file on the fly?

[mcwitt]

It seems like CMAKE will create this file on the fly?

Correct; the original reason to check it in was so it can be used by the mypy check during linting. But since we're now building the extension during linting anyway to verify that the stubs are up to date, seems like we might be able to avoid checking it in. I'll look into this.

[mcwitt]

Another benefit of checking it in is to get typechecking and completions to work when developing the Python code on non-Linux machines. Not sure how common of a use case that is, though.. It seems like most of us are either developing locally on Linux, or running a remote IDE server on a Linux box?

@maxentile @proteneer @jkausrelay
do you know if having the stub file locally would be useful for your dev setups?

[maxentile]

The changes this suggests mostly look like improvements to me!

[maxentile]

Nice -- didn't realize this was duplicated!

[maxentile]

Mild surprise: I thought this was often an array

[mcwitt]

Good call; I don't think we have any cases of passing an array in the timemachine code, but there are plenty of examples in the tests (we don't currently run mypy on tests; this seemed like maybe too much overhead).

I think a good general approach might be to use the most general type that has the capabilities exercised by the implementation. In this case, we use __getitem__ and __len__, which suggests Sequence might be appropriate. numpy.ndarray does not implement one of the required methods of Sequence (.index()), but we don't use that anyway. So, an appropriate type here might be something like Union[Sequence[float], NDArray].

[mcwitt]

Changed this in ff5c69b

[maxentile]

Tangent / curiosity: It's often convenient to use array containers and python containers interchangeably in this sort of code.

Is there a more generic type for this situation? (should accept: int array of shape (n_pairs, 2), a list of 2-tuples, etc... Initial thought is something like Union[Array, List[Tuple[int, int]]], but that seems counterproductive since more verbose, and since array typing is less precise)

[mcwitt]

Initial thought is something like Union[Array, List[Tuple[int, int]]]

I think this is a reasonable compromise -- we lose some power to ensure that the input is the correct shape, but this still contains useful information (if not an array, we know the input is a list of 2-tuples)

[mcwitt]

In this case, I think the most general type compatible with the implementation might be something like

validate_map(n_nodes: int, relative_inds: Iterable[Tuple[int, int]], absolute_inds: Collection[int]) -> bool

For example, this would work with an array for relative_inds and a tuple for absolute_inds

[mcwitt]

(I think this sort of generalization might be best done as an iterative process; trying to come up with the most general types initially is a fair amount of work and error-prone. For example, I might make the above update in a PR where I add some code that passes an array to relative_inds)

[maxentile]

Nice: Iterable[Tuple[int, int]] seems both generic and precise

(I think this sort of generalization might be best done as an iterative process; trying to come up with the most general types initially is a fair amount of work and error-prone.

Agreed!