ENH: Adding an empty __init__ method to distribution classes

[thomasaarholt]

Before

Currently, IDE's are unable to display the correct variable and type signature of pymc's distributions. This has been discussed before, see #6083.

This is the current on-hover behaviour of the main branch on pm.Normal:

Now that the distributions .dist methods are well-typed (see #6635 and #6937, the latter just needs review and merge), I propose that we simply copy the method type signature over to an "empty" (returns None) __init__ method for each distribution.

This should let us have the correct type signature (no more *args, a huge benefit) with hinting (a big benefit) with a small cost of some code duplication.

After

Let me provide some evidence that the above method will work (here in VSCode).

Code change

Here I make only the following code change to pymc's pm.Normal class:

Result

On hover of the Normal class, we see a much nicer function definition:

Context for the issue:

• There is a discussion in Add distribution parameters as named positional arguments #6083 that it would be a good idea to refactor the distributions from classes into functions. I agree with that, but since there has been no effort to do so in the last 1.5 years, I suggest that this issue's approach will suffice until such an effort can be made.

I think it might be useful in taking variables out of kwargs and adding them directly in the __init__, however, like shape, dims and observed. Please come with thoughts and suggestions in the discussion below. I have never seen total_size, initval or transform used, but that doesn't mean that they shouldn't be added.

• such changes, as well as a merge of 6937 should be done before creating the __init__ methods, so as to reduce duplicate work.

I am happy to make a PR with the proposed changes, perhaps focusing on the most common distributions first, to keep PRs small and reviewable.

This is something that has been bothering me for a while, and I am quite keen to see it added.

[ricardoV94]

Can that be automated? The MetaClass is already doing a lot of magic so maybe it can also add docs for the dummy __init__. I wouldn't want the signature to be duplicated manually, as it will eventually diverge/ be a source of confusion for developers.

I did something like this in DistributionMeta.__new__():

        clsdict.setdefault("__init__", clsdict.get("dist", None))

It shows up in help(pm.Normal) but not in the tooltip in PyCharm

[ricardoV94]

Since the solution here is specific to the dummy classes, I tried to push the discussion in #5308 with a concrete example in: #5308 (comment)

[thomasaarholt]

Hi Ricardo! Thanks for replying so quick to this.

In general, the type checkers (pylance/pyright for vscode, jedi for jupyter lab, and this one in pycharm, which also generate the tooltip that we see when hovering over a class or function, are static type checkers.

Unfortunately, that means that any dynamic behaviour, like setting the call signature of one function to equal another one, does not work.

So this approach does indeed require duplicating the call signature. Is that that bad though? I don't have the impression that the majority of these functions change particularly frequently. We could even try to reduce the kwargs namespace using TypedDict (pre-3.12 support exists through typing_extensions). I wrote up a little showcase in this gist. See the comment below the code.

I do like your suggestion in the discussion though, will respond there in a bit.