Skip to content

Commit

Permalink
Fix MyPy's view of _ModuleMeta
Browse files Browse the repository at this point in the history
  • Loading branch information
NeilGirdhar authored and patrick-kidger committed Aug 31, 2024
1 parent 2508192 commit 3cdeef8
Showing 1 changed file with 7 additions and 3 deletions.
10 changes: 7 additions & 3 deletions equinox/_module.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 = {
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 3cdeef8

Please sign in to comment.