Fix ruff and types

AlignmentResearch · Oct 11, 2023 · 1e7cf77 · 1e7cf77
1 parent 57dfc30
commit 1e7cf77
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 14 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -51,7 +51,7 @@ jobs:
           command: ruff .
       - run:
           name: Typecheck (mypy)
-          command: mypy stable_baselines3/common/pytree_dataclass.py
+          command: mypy stable_baselines3/common/pytree_dataclass.py tests/test_pytree_dataclass.py  # TODO: remove, in PR#4.
   pytype:
     docker:
       - image: ghcr.io/alignmentresearch/learned-planners:<< pipeline.parameters.docker_img_version >>
@@ -62,7 +62,7 @@ jobs:
     working_directory: /workspace/third_party/stable-baselines3
     steps:
       - checkout
-      - run: pytype -j 4 stable_baselines3/common/pytree_dataclass.py
+      - run: pytype -j 4  stable_baselines3/common/pytree_dataclass.py tests/test_pytree_dataclass.py  # TODO: remove, in PR#4.
   py-tests:
     docker:
       - image: ghcr.io/alignmentresearch/learned-planners:<< pipeline.parameters.docker_img_version >>

diff --git a/tests/test_pytree_dataclass.py b/tests/test_pytree_dataclass.py
@@ -37,7 +37,7 @@ class D(ParentPyTreeClass):
 
     with pytest.raises(TypeError, match="You should not specify frozen= for descendants"):
 
-        class D(ParentPyTreeClass, frozen=True):  # type: ignore
+        class D(ParentPyTreeClass, frozen=True):  # type: ignore  # noqa:F811
             a: int
 
 
@@ -53,17 +53,17 @@ class D(ptd._PyTreeDataclassBase, frozen=frozen):  # type: ignore
 
     with pytest.raises(TypeError):
 
-        class D(metaclass=ptd._PyTreeDataclassMeta, frozen=frozen):  # type: ignore
+        class D(metaclass=ptd._PyTreeDataclassMeta, frozen=frozen):  # type: ignore  # noqa: F811
             pass
 
     with pytest.raises(TypeError, match="[^ ]* dataclass .* should inherit"):
 
-        class D(ptd._PyTreeDataclassBase):  # type: ignore
+        class D(ptd._PyTreeDataclassBase):  # type: ignore  # noqa: F811
             pass
 
     with pytest.raises(TypeError, match="[^ ]* dataclass .* should inherit"):
 
-        class D(metaclass=ptd._PyTreeDataclassMeta):  # type: ignore
+        class D(metaclass=ptd._PyTreeDataclassMeta):  # type: ignore  # noqa: F811
             pass
 
     # Then try to copy each of the reserved names:
@@ -96,17 +96,17 @@ class FrozenPyTreeDataclass(ptd._PyTreeDataclassBase, frozen=frozen):  # type: i
 
     with pytest.raises(TypeError):
 
-        class FrozenPyTreeDataclass(metaclass=ptd._PyTreeDataclassMeta, frozen=frozen):  # type: ignore
+        class FrozenPyTreeDataclass(metaclass=ptd._PyTreeDataclassMeta, frozen=frozen):  # type: ignore  # noqa: F811
             pass
 
     with pytest.raises(TypeError, match="You cannot have another class named"):
 
-        class FrozenPyTreeDataclass(ptd._PyTreeDataclassBase):  # type: ignore
+        class FrozenPyTreeDataclass(ptd._PyTreeDataclassBase):  # type: ignore  # noqa: F811
             pass
 
     with pytest.raises(TypeError, match="You cannot have another class named"):
 
-        class FrozenPyTreeDataclass(metaclass=ptd._PyTreeDataclassMeta):  # type: ignore
+        class FrozenPyTreeDataclass(metaclass=ptd._PyTreeDataclassMeta):  # type: ignore  # noqa: F811
             pass
 
     ## MutablePyTreeDataclass
@@ -117,17 +117,17 @@ class MutablePyTreeDataclass(ptd._PyTreeDataclassBase, frozen=frozen):  # type:
 
     with pytest.raises(TypeError):
 
-        class MutablePyTreeDataclass(metaclass=ptd._PyTreeDataclassMeta, frozen=frozen):  # type: ignore
+        class MutablePyTreeDataclass(metaclass=ptd._PyTreeDataclassMeta, frozen=frozen):  # type: ignore  # noqa:F811
             pass
 
     with pytest.raises(TypeError, match="You cannot have another class named"):
 
-        class MutablePyTreeDataclass(ptd._PyTreeDataclassBase):  # type: ignore
+        class MutablePyTreeDataclass(ptd._PyTreeDataclassBase):  # type: ignore  # noqa:F811
             pass
 
     with pytest.raises(TypeError, match="You cannot have another class named"):
 
-        class MutablePyTreeDataclass(metaclass=ptd._PyTreeDataclassMeta):  # type: ignore
+        class MutablePyTreeDataclass(metaclass=ptd._PyTreeDataclassMeta):  # type: ignore  # noqa:F811
             pass
 
 
@@ -166,15 +166,16 @@ def test_tree_flatten() -> None:
     class A(ptd.FrozenPyTreeDataclass):
         a: Optional[int]
 
-    flat, _ = ptd.tree_flatten((A(3), A(None), {"a": A(4)}))
+    flat, _ = ptd.tree_flatten((A(3), A(None), {"a": A(4)}))  # type: ignore
     assert flat == [3, 4]
 
 
 def test_tree_map() -> None:
     class A(ptd.FrozenPyTreeDataclass):
         a: Optional[int]
 
-    assert ptd.tree_map(lambda x: x * 2, ([2, 3], 4, A(5), None, {"a": 6})) == ([4, 6], 8, A(10), None, {"a": 12})  # type: ignore
+    out = ptd.tree_map(lambda x: x * 2, ([2, 3], 4, A(5), None, {"a": 6}))  # type: ignore
+    assert out == ([4, 6], 8, A(10), None, {"a": 12})
 
 
 def test_tree_empty() -> None: