How to type hint with InferenceData

[williambdean]

How can I avoid this?

Aren't the groups "defined" or restricted to a small set?

Might be related to #1581 but I interpreted that as internal type hinting

[tomicapretto]

What I've done in the past is idata["prior"], etc. But I don't know if that's a real solution.

[amaloney]

Thanks @tomicapretto for the suggestion. I agree this is the "best" method when using ArviZ and what looks to me is pyright in neovim @wd60622.

The reason why one would get these type errors is because the InferenceData object does not explicitly define anywhere in the code the name prior, posterior, etc. Instead, what we do is inherit the Mapping object from Python into the InferenceData object, with keys as strings and values as xarray Dataset objects, see below.

arviz/arviz/data/inference_data.py

Line 94 in 0fc1117

class InferenceData(Mapping[str, xr.Dataset]):

Allowing you to use "dot" notation is built into the InferenceData object, which is why you can use it, but it will give you the type errors you see. To remove those errors, I would also suggest doing the following.

def function(idata: az.InferenceData) -> None:
    idata["prior"]
    idata["posterior"]
    idata["posterior_predictive"]
    idata["sample_stats"]
    idata["log_likelihood"]

Then you will not get the type errors you see, since you are accessing the InferenceData object like what its inherited type is, a Mapping, or dictionary.

[williambdean]

I personally dont use brackets. Seems imprecise.

Does this warn with idata["unlikely_key"]?

[amaloney]

No, but neither would a dictionary, which is what the InferenceData object is based on. This is a deep issue with how Python objects are created, and I do not see an immediate answer. I have seen this type confusion before with InferenceData objects and other users, so you are in good company.

Let me think about it more, since there are ways to allow for this notation in Python. Take for example the following where I've converted the InferenceData object to a SimpleNamespace object. Doing this has many downstream implications, but you no longer have a type error. Doing this is not recommended without further testing and analysis