From 3cdeef80c59a1092c0205319bf7fbadd9befdcd0 Mon Sep 17 00:00:00 2001 From: Neil Girdhar Date: Thu, 29 Aug 2024 16:19:45 -0400 Subject: [PATCH] Fix MyPy's view of _ModuleMeta --- equinox/_module.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/equinox/_module.py b/equinox/_module.py index de5b7f1f..3c726b80 100644 --- a/equinox/_module.py +++ b/equinox/_module.py @@ -176,7 +176,7 @@ class AbstractFoo(eqx.Module, strict=True): # Inherits from ABCMeta to support `eqx.{AbstractVar, AbstractClassVar}` and # `abc.abstractmethod`. -class _ModuleMeta(ABCMeta): # pyright: ignore +class _ActualModuleMeta(ABCMeta): # pyright: ignore # This method is called whenever you definite a module: `class Foo(eqx.Module): ...` def __new__( mcs, @@ -569,7 +569,7 @@ def __call__(cls, *args, **kwargs): post_init = getattr(cls, "__post_init__", None) initable_cls = _make_initable(cls, cls.__init__, post_init, wraps=False) # [Step 2] Instantiate the class as normal. - self = super(_ModuleMeta, initable_cls).__call__(*args, **kwargs) + self = super(_ActualModuleMeta, initable_cls).__call__(*args, **kwargs) assert not _is_abstract(cls) # [Step 3] Check that all fields are occupied. missing_names = { @@ -638,6 +638,8 @@ def __setattr__(cls, item, value): class _ModuleMeta(abc.ABCMeta): __abstractvars__: frozenset[str] __abstractclassvars__: frozenset[str] +else: + _ModuleMeta = _ActualModuleMeta def _is_special_form(cls): @@ -791,7 +793,9 @@ def __call__(self, ...): @ft.lru_cache(maxsize=128) -def _make_initable(cls: _ModuleMeta, init, post_init, wraps: bool) -> _ModuleMeta: +def _make_initable( + cls: _ActualModuleMeta, init, post_init, wraps: bool +) -> _ActualModuleMeta: # Used as part of the key. Don't cache if these have changed. # In practice, monkey-patching these on the class -- after you've already # instantiated it somewhere! -- is an *ahem*, adventurous, thing to do. But never