feat(dataclasses): allow creating dataclass types in the global namespace (#212)

Author: XuehaiPan

CHANGELOG.md

@@ -18,6 +18,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 - Enforce naming convention of packages with singular and plural: `optree.{accessor,integration}` -> `optree.{accessors,integrations}` by [@XuehaiPan](https://github.com/XuehaiPan) in [#209](https://github.com/metaopt/optree/pull/209).
+- Allow creating dataclass types in the global namespace by [@XuehaiPan](https://github.com/XuehaiPan) in [#212](https://github.com/metaopt/optree/pull/212).
 
 ### Fixed
 

optree/dataclasses.py

@@ -302,7 +302,7 @@ def decorator(cls: _TypeT) -> _TypeT:
     if namespace is not GLOBAL_NAMESPACE and not isinstance(namespace, str):
         raise TypeError(f'The namespace must be a string, got {namespace!r}.')
     if namespace == '':
-        raise ValueError('The namespace cannot be an empty string.')
+        namespace = GLOBAL_NAMESPACE
 
     cls = dataclasses.dataclass(cls, **kwargs)  # type: ignore[assignment]
 
@@ -412,7 +412,7 @@ def make_dataclass(  # type: ignore[no-redef] # noqa: C901,D417
     if namespace is not GLOBAL_NAMESPACE and not isinstance(namespace, str):
         raise TypeError(f'The namespace must be a string, got {namespace!r}.')
     if namespace == '':
-        raise ValueError('The namespace cannot be an empty string.')
+        namespace = GLOBAL_NAMESPACE
 
     dataclass_kwargs = {
         'init': init,

tests/test_dataclasses.py

@@ -368,30 +368,41 @@ class Foo1:
             x: int
             y: float
 
-    with pytest.raises(ValueError, match=re.escape('The namespace cannot be an empty string.')):
+    @optree.dataclasses.dataclass(namespace='')
+    class Foo2:
+        x: int
+        y: float
 
-        @optree.dataclasses.dataclass(namespace='')
-        class Foo2:
-            x: int
-            y: float
+    foo = Foo2(1, 2.0)
+    accessors, leaves, treespec = optree.tree_flatten_with_accessor(foo)
+    assert optree.tree_unflatten(treespec, leaves) == foo
+    assert accessors == [
+        optree.PyTreeAccessor((optree.DataclassEntry('x', Foo2, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('y', Foo2, optree.PyTreeKind.CUSTOM),)),
+    ]
+    assert [a(foo) for a in accessors] == [1, 2.0]
+    assert leaves == [1, 2.0]
+    assert treespec.namespace == ''
+    assert treespec.kind == optree.PyTreeKind.CUSTOM
+    assert treespec.type is Foo2
 
     @optree.dataclasses.dataclass(namespace=GLOBAL_NAMESPACE)
-    class Foo:
+    class Foo3:
         x: int
         y: float
 
-    foo = Foo(1, 2.0)
+    foo = Foo3(1, 2.0)
     accessors, leaves, treespec = optree.tree_flatten_with_accessor(foo)
     assert optree.tree_unflatten(treespec, leaves) == foo
     assert accessors == [
-        optree.PyTreeAccessor((optree.DataclassEntry('x', Foo, optree.PyTreeKind.CUSTOM),)),
-        optree.PyTreeAccessor((optree.DataclassEntry('y', Foo, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('x', Foo3, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('y', Foo3, optree.PyTreeKind.CUSTOM),)),
     ]
     assert [a(foo) for a in accessors] == [1, 2.0]
     assert leaves == [1, 2.0]
     assert treespec.namespace == ''
     assert treespec.kind == optree.PyTreeKind.CUSTOM
-    assert treespec.type is Foo
+    assert treespec.type is Foo3
 
 
 def test_make_dataclass_future_parameters():
@@ -684,30 +695,41 @@ def test_make_dataclass_with_invalid_namespace():
     with pytest.raises(TypeError, match='The namespace must be a string'):
         optree.dataclasses.make_dataclass('Foo1', ['x', ('y', int), ('z', float, 0.0)], namespace=1)
 
-    with pytest.raises(ValueError, match=re.escape('The namespace cannot be an empty string.')):
-        optree.dataclasses.make_dataclass(
-            'Foo2',
-            ['x', ('y', int), ('z', float, 0.0)],
-            namespace='',
-        )
+    Foo2 = optree.dataclasses.make_dataclass(  # noqa: N806
+        'Foo2',
+        [('x', int), ('y', float)],
+        namespace='',
+    )
+    foo = Foo2(1, 2.0)
+    accessors, leaves, treespec = optree.tree_flatten_with_accessor(foo)
+    assert optree.tree_unflatten(treespec, leaves) == foo
+    assert accessors == [
+        optree.PyTreeAccessor((optree.DataclassEntry('x', Foo2, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('y', Foo2, optree.PyTreeKind.CUSTOM),)),
+    ]
+    assert [a(foo) for a in accessors] == [1, 2.0]
+    assert leaves == [1, 2.0]
+    assert treespec.namespace == ''
+    assert treespec.kind == optree.PyTreeKind.CUSTOM
+    assert treespec.type is Foo2
 
-    Foo = optree.dataclasses.make_dataclass(  # noqa: N806
-        'Foo',
+    Foo3 = optree.dataclasses.make_dataclass(  # noqa: N806
+        'Foo3',
         [('x', int), ('y', float)],
         namespace=GLOBAL_NAMESPACE,
     )
-    foo = Foo(1, 2.0)
+    foo = Foo3(1, 2.0)
     accessors, leaves, treespec = optree.tree_flatten_with_accessor(foo)
     assert optree.tree_unflatten(treespec, leaves) == foo
     assert accessors == [
-        optree.PyTreeAccessor((optree.DataclassEntry('x', Foo, optree.PyTreeKind.CUSTOM),)),
-        optree.PyTreeAccessor((optree.DataclassEntry('y', Foo, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('x', Foo3, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('y', Foo3, optree.PyTreeKind.CUSTOM),)),
     ]
     assert [a(foo) for a in accessors] == [1, 2.0]
     assert leaves == [1, 2.0]
     assert treespec.namespace == ''
     assert treespec.kind == optree.PyTreeKind.CUSTOM
-    assert treespec.type is Foo
+    assert treespec.type is Foo3
 
     with pytest.raises(TypeError, match='The namespace must be a string'):
         optree.dataclasses.make_dataclass(
@@ -717,29 +739,40 @@ def test_make_dataclass_with_invalid_namespace():
             namespace=None,
         )
 
-    with pytest.raises(ValueError, match=re.escape('The namespace cannot be an empty string.')):
-        optree.dataclasses.make_dataclass(
-            'Foo2',
-            ['x', ('y', int), ('z', float, 0.0)],
-            ns='',
-            namespace={},
-        )
+    Bar2 = optree.dataclasses.make_dataclass(  # noqa: N806
+        'Bar2',
+        [('x', int), ('y', float)],
+        ns='',
+        namespace={},
+    )
+    bar = Bar2(1, 2.0)
+    accessors, leaves, treespec = optree.tree_flatten_with_accessor(bar)
+    assert optree.tree_unflatten(treespec, leaves) == bar
+    assert accessors == [
+        optree.PyTreeAccessor((optree.DataclassEntry('x', Bar2, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('y', Bar2, optree.PyTreeKind.CUSTOM),)),
+    ]
+    assert [a(bar) for a in accessors] == [1, 2.0]
+    assert leaves == [1, 2.0]
+    assert treespec.namespace == ''
+    assert treespec.kind == optree.PyTreeKind.CUSTOM
+    assert treespec.type is Bar2
 
-    Bar = optree.dataclasses.make_dataclass(  # noqa: N806
-        'Bar',
+    Bar3 = optree.dataclasses.make_dataclass(  # noqa: N806
+        'Bar3',
         [('x', int), ('y', float)],
         ns=GLOBAL_NAMESPACE,
         namespace={},
     )
-    bar = Bar(1, 2.0)
+    bar = Bar3(1, 2.0)
     accessors, leaves, treespec = optree.tree_flatten_with_accessor(bar)
     assert optree.tree_unflatten(treespec, leaves) == bar
     assert accessors == [
-        optree.PyTreeAccessor((optree.DataclassEntry('x', Bar, optree.PyTreeKind.CUSTOM),)),
-        optree.PyTreeAccessor((optree.DataclassEntry('y', Bar, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('x', Bar3, optree.PyTreeKind.CUSTOM),)),
+        optree.PyTreeAccessor((optree.DataclassEntry('y', Bar3, optree.PyTreeKind.CUSTOM),)),
     ]
     assert [a(bar) for a in accessors] == [1, 2.0]
     assert leaves == [1, 2.0]
     assert treespec.namespace == ''
     assert treespec.kind == optree.PyTreeKind.CUSTOM
-    assert treespec.type is Bar
+    assert treespec.type is Bar3