cgarciae / simple-pytree

A dead simple Python package for creating custom JAX pytree objects

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Generic typing issue?

daniel-dodd opened this issue · comments

Hi @cgarciae,

Thank you for creating such a neat library. Have quite enjoyed playing around with it.

I have potentially run into an issue with generic typing:

# version == 0.1.4

from simple_pytree import Pytree

import typing as tp

T = tp.TypeVar("T")

class MyClass(Pytree, tp.Generic[T]):
    def __init__(self, x: T):
        self.x = x

MyClass[int]
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
Cell In[1], line 14
     11     def __repr__(self) -> str:
     12         return f"MyClass({self.x})"
---> 14 MyClass[int]

File [~/miniconda3/lib/python3.10/typing.py:309](https://file+.vscode-resource.vscode-cdn.net/Users/danieldodd/Desktop/~/miniconda3/lib/python3.10/typing.py:309), in _tp_cache..decorator..inner(*args, **kwds)
    306 @functools.wraps(func)
    307 def inner(*args, **kwds):
    308     try:
--> 309         return cached(*args, **kwds)
    310     except TypeError:
    311         pass  # All real errors (not unhashable args) are raised below.

File [~/miniconda3/lib/python3.10/typing.py:1342](https://file+.vscode-resource.vscode-cdn.net/Users/danieldodd/Desktop/~/miniconda3/lib/python3.10/typing.py:1342), in Generic.__class_getitem__(cls, params)
   1338         raise TypeError(
   1339             f"Parameters to {cls.__name__}[...] must all be unique")
   1340 else:
   1341     # Subscripting a regular Generic subclass.
-> 1342     if any(isinstance(t, ParamSpec) for t in cls.__parameters__):
   1343         params = _prepare_paramspec_params(cls, params)
   1344     else:

AttributeError: type object 'MyClass' has no attribute '__parameters__'

Would appreciate any insight to what might be causing this issue, if you know the answer. Otherwise happy to investigate/support you with a PR, if this not expected behaviour.

Cheers,
Dan :)

Hey @daniel-dodd, thanks for the kind words :)

I've opened a PR to fix this, will publish a new version soon. In the meanwhile, you can fix the problem by simply changing the order of the parent classes:

class MyClass(tp.Generic[T], Pytree):

Fixed on 0.1.5