Add type hints to EpochsArray constructor and BaseRaw.get_data()

[hoechenberger]

This is WIP and I'm trying out some things while also simultaneously attempting not to go insane.

Whether we like it or not, as you know, most IDEs these days rely on static code analysis to assist writing code and offering tab completions etc.

So to provide a better UX, I started experimenting with type annotations in MNE-Python.

Now there are several design patterns we have in our code that make me pull out my hair – not because I necessarily think they're inherently bad, not all of them; but because some are extremely difficult or impossible to properly annotate such that static analysis tools will be able to infer the correct types and we don't have to bend over backwards

One of such recurring patterns, for example, is changing the type of a return value depending on a function parameter.

Let's look at the following code:

# %%

import numpy as np

import mne

sample_dir = mne.datasets.sample.data_path()
sample_fname = sample_dir / 'MEG' / 'sample' / 'sample_audvis_raw.fif'

raw = mne.io.read_raw_fif(sample_fname, preload=True)
raw.crop(tmax=60)

data = raw.get_data()
assert isinstance(data, np.ndarray)
raw_array = mne.io.RawArray(data=data, info=raw.info)

raw.get_data() has a switch, return_times, which will turn the return value from NDArray into a tuple[NDArray, NDArray].

This means that which type will actually be returned can only be determined at runtime, except if one uses @overloads, which we figured are a bit … heavy? Let's call it heavy :)

But without adding an overload to get_data() to handle the two cases – return_times=False -> NDArray, return_times=True -> tuple, the user is forced to perform a runtime type check. Without the assert in my MWE above, mne.io.RawArray(data=data, …) would (correctly!) get red squiggly lines, because Pyright cannot know whether it's dealing with an array or a tuple of arrays here, while the constructor clearly demands an array.

The only solution to that problem that wouldn't require a change or amendment to our API is adding overloads to get_data().

But since previous consensus was that we're not super happy about this approach, I'm right now a little bit lost and demotivated on how to best proceed. I want to provide an improved UX but I fully understand the reservations towards overloads.

That said … This PR includes some improvements to EpochsArray and BaseRaw.get_data() that already improve the UX a bit but I'm not happy about the unclear return type of get_data(). Sigh…

[cbrnr]

I feel you! In certain cases overloads are totally fine, maybe even in this one (or rather, in all cases where return types vary as a function of arguments). If we don't want that, it's best to separate into two or more functions with stable return types.