Fix some type hinting to help with migrating Distribution

[thomasaarholt]

Description

This is a small PR. I've added and fixed type hinting to a few lines of code. The idea is to make it easier to migrate the Distribution class to a function-based approach as discussed in the related issue.

Related Issue

• Related to Add distribution parameters as named positional arguments #6083

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pymc--7484.org.readthedocs.build/en/7484/

[codecov]

Codecov Report

Attention: Patch coverage is 90.00000% with 1 line in your changes missing coverage. Please review.

Project coverage is 92.16%. Comparing base (d596afb) to head (157c29a).

Files with missing lines	Patch %	Lines
pymc/distributions/distribution.py	80.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7484      +/-   ##
==========================================
- Coverage   92.16%   92.16%   -0.01%     
==========================================
  Files         103      103              
  Lines       17200    17203       +3     
==========================================
+ Hits        15853    15855       +2     
- Misses       1347     1348       +1     

Files with missing lines	Coverage Δ
pymc/distributions/continuous.py	97.92% <100.00%> (+<0.01%)	⬆️
pymc/distributions/distribution.py	91.66% <80.00%> (-0.23%)	⬇️

[ricardoV94]

mypy is not happy with the changes? https://github.com/pymc-devs/pymc/actions/runs/10613016885/job/29415940076?pr=7484

[thomasaarholt]

Trying a solution, see below.

Yes. That's this commit, where I fixed the incorrect type hint syntax.

All three reported errors are the same:

Incompatible types in assignment 
(expression has type "Callable[[Any, VarArg(Any), DefaultNamedArg(Any, 'size')], Any]", 
base class "Distribution" defined the type as "RandomVariable | SymbolicRandomVariable")

and are due to Mixture, Censored and Discrete distributions having rv_op = <a method>, while e.g. Normal has rv_op = <a class>.

class Censored(Distribution):
    rv_type = CensoredRV
    rv_op = CensoredRV.rv_op # <-- a method

class Normal(Distribution):
    rv_op = NormalRV() # <-- an instance of NormalRV, which is a subclass of RandomVariable

Solution
Typing the above, I realize that since we are really wanting to type this by behaviour rather than state. So I've tried typing it as Callable[..., TensorVariable] instead, which is how it always behaves (I think).

[ricardoV94]

sounds correct

[thomasaarholt]

I am not going to worry about the error in func_utils.py. It's in find_constrained_prior, which is a) deprecated and b) not typed correctly (first argument should be type[Distribution]).

I am currently testing if I can forego the = None on the rv_op and rv_type completely. I don't think these are necessary.

[thomasaarholt]

The failing test is unrelated I guess? @ricardoV94

I think I'm finished with this one 😊

[ricardoV94]

I am not going to worry about the error in func_utils.py. It's in find_constrained_prior, which is a) deprecated and b) not typed correctly (first argument should be type[Distribution]).

I am currently testing if I can forego the = None on the rv_op and rv_type completely. I don't think these are necessary.

Then can you remove the type hints, we can't merge with a known failing CI or it will be a source of confusion for everyone else

[thomasaarholt]

I can remove the type hint for class attribute rv_op, but then I either need to leave it = None (which the passing tests just established is unnecessary), replace it (with Any, probably, and a comment explaining what I've learnt about it), or remove the entire attribute, since the following is not legal Python:

class Foo:
    bar

[thomasaarholt]

@ricardoV94 mind taking a look?

[ricardoV94]

Fun debug statement left :)

Suggested change

	if class_change_dist_size is not None:
	raise ValueError("HAHAHA")

[thomasaarholt]

Hehe! That's what I get for not adding my changes line by line for this particular commit!

[ricardoV94]

Aren't there distributions with only one of the bounds (lower, upper)?

[thomasaarholt]

You're right, there are! I was wondering why mypy didn't flag it, but of course it was originally an error as list[int] doesn't allow (3, None)! Fixing this

[thomasaarholt]

The failing tests stem from this commit, not this PR.

[ricardoV94]

@thomasaarholt should be fixed by #7529

[ricardoV94]

Thanks @thomasaarholt