Module bug

[ASEM000]

Hello patrick,

I think the following is a bug in Module

import equinox as eqx

class ParentTree(eqx.Module):
    def __init_subclass__(klass: type, foo: bool = False):
        if foo:
            klass.foo = True

class Tree(ParentTree, foo=True):
    def __init__(self, a: int):
        self.a = a

Tree(a=1).foo
# _ModuleMeta.__new__() got an unexpected keyword argument 'foo'

if i omit foo a different error comes up:

class Tree(ParentTree):
    def __init__(self, a: int):
        self.a = a

Tree(a=1)
# AttributeError: Cannot set attribute a

In comparison to pytreeclass or simple_pytree, this works as expected:

import pytreeclass as pytc
class ParentTree(pytc.TreeClass):
    def __init_subclass__(klass: type, foo: bool = False):
        if foo:
            klass.foo = True
class Tree(ParentTree, foo=True):
    def __init__(self, a: int):
        self.a = a
Tree(a=1).foo
# True

from simple_pytree import Pytree
class ParentTree(Pytree):
    def __init_subclass__(klass: type, foo: bool = False):
        if foo:
            klass.foo = True
class Tree(ParentTree, foo=True):
    def __init__(self, a: int):
        self.a = a
Tree(a=1).foo
# True

TheModule logic became more complex since the last time i checked, so Im not sure the reason of the two errors. But one pointer is that maybe propagate kwargs in module meta new.

[patrick-kidger]

Regarding foo=True: thanks! This'll be fixed by #482.

Regarding omitting foo: Equinox's syntax involves explicitly declaring all attributes. You want

class Tree(ParentTree):
    a: int  # this bit is new

    def __init__(self, a: int):
        self.a = a

although also note that the default __init__ method matches the one explicity provided, so that can be skipped. Both the attribute declaration and the default __init__ are the same behaviour as with any dataclass.

[ASEM000]

This was quick :D

[ASEM000]

Regarding omitting foo: Equinox's syntax involves explicitly declaring all attributes.

My bad.

[patrick-kidger]

Closing! Fixed in #482, and this will appear in the upcoming v0.11.0 release.