ruff formatting

docs/developer_guide/example_docstrings.py

@@ -8,7 +8,7 @@ def example_function_args(
     param2: str | int,
     optional_param: list[str] | None = None,
     *args: float | str,
-    **kwargs: Dict[str, Any]
+    **kwargs: Dict[str, Any],
 ) -> tuple[bool, list[str]]:
     """Example function with PEP 484 type annotations and PEP 563 future annotations.
 

medmodels/_medmodels.pyi

@@ -1,4 +1,3 @@
-
 from enum import Enum
 from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Union
 

medmodels/medrecord/medrecord.py

@@ -776,9 +776,9 @@ def add_edges(
             edges
         ):
             return self.add_edges_pandas(edges, group)
-        if is_polars_edge_dataframe_input(
+        if is_polars_edge_dataframe_input(edges) or is_polars_edge_dataframe_input_list(
             edges
-        ) or is_polars_edge_dataframe_input_list(edges):
+        ):
             return self.add_edges_polars(edges, group)
         if is_edge_tuple(edges):
             edges = [edges]

medmodels/medrecord/schema.py

@@ -252,8 +252,12 @@ class GroupSchema:
     def __init__(
         self,
         *,
-        nodes: Optional[Dict[MedRecordAttribute, Union[DataType, Tuple[DataType, AttributeType]]]] = None,
-        edges: Optional[Dict[MedRecordAttribute, Union[DataType, Tuple[DataType, AttributeType]]]] = None,
+        nodes: Optional[
+            Dict[MedRecordAttribute, Union[DataType, Tuple[DataType, AttributeType]]]
+        ] = None,
+        edges: Optional[
+            Dict[MedRecordAttribute, Union[DataType, Tuple[DataType, AttributeType]]]
+        ] = None,
         strict: bool = False,
     ) -> None:
         """Initializes a new instance of GroupSchema.

medmodels/medrecord/tests/test_datatype.py

@@ -107,7 +107,10 @@ def test_union(self) -> None:
 
         assert str(union) == "Union(String, Union(Int, Bool))"
 
-        assert union.__repr__() == "DataType.Union(DataType.String, DataType.Union(DataType.Int, DataType.Bool))"
+        assert (
+            union.__repr__()
+            == "DataType.Union(DataType.String, DataType.Union(DataType.Int, DataType.Bool))"
+        )
 
         assert mr.Union(mr.String(), mr.Int()) == mr.Union(mr.String(), mr.Int())
         assert mr.Union(mr.String(), mr.Int()) != mr.Union(mr.Int(), mr.String())

medmodels/medrecord/tests/test_medrecord.py

@@ -327,7 +327,10 @@ def test_group(self) -> None:
 
         assert medrecord.group("1") == {"nodes": ["0"], "edges": [0]}
 
-        assert medrecord.group(["0", "1"]) == {"0": {"nodes": [], "edges": []}, "1": {"nodes": ["0"], "edges": [0]}}
+        assert medrecord.group(["0", "1"]) == {
+            "0": {"nodes": [], "edges": []},
+            "1": {"nodes": ["0"], "edges": [0]},
+        }
 
     def test_invalid_group(self) -> None:
         medrecord = create_medrecord()
@@ -351,11 +354,17 @@ def test_outgoing_edges(self) -> None:
 
         edges = medrecord.outgoing_edges(["0", "1"])
 
-        assert {key: sorted(value) for key, value in edges.items()} == {"0": sorted([0, 3]), "1": [1, 2]}
+        assert {key: sorted(value) for key, value in edges.items()} == {
+            "0": sorted([0, 3]),
+            "1": [1, 2],
+        }
 
         edges = medrecord.outgoing_edges(node_select().index().is_in(["0", "1"]))
 
-        assert {key: sorted(value) for key, value in edges.items()} == {"0": sorted([0, 3]), "1": [1, 2]}
+        assert {key: sorted(value) for key, value in edges.items()} == {
+            "0": sorted([0, 3]),
+            "1": [1, 2],
+        }
 
     def test_invalid_outgoing_edges(self) -> None:
         medrecord = create_medrecord()
@@ -1383,7 +1392,10 @@ def test_groups_of_node(self) -> None:
 
         assert medrecord.groups_of_node(["0", "1"]) == {"0": ["0"], "1": ["0"]}
 
-        assert medrecord.groups_of_node(node_select().index().is_in(["0", "1"])) == {"0": ["0"], "1": ["0"]}
+        assert medrecord.groups_of_node(node_select().index().is_in(["0", "1"])) == {
+            "0": ["0"],
+            "1": ["0"],
+        }
 
     def test_invalid_groups_of_node(self) -> None:
         medrecord = create_medrecord()
@@ -1405,7 +1417,10 @@ def test_groups_of_edge(self) -> None:
 
         assert medrecord.groups_of_edge([0, 1]) == {0: ["0"], 1: ["0"]}
 
-        assert medrecord.groups_of_edge(edge_select().index().is_in([0, 1])) == {0: ["0"], 1: ["0"]}
+        assert medrecord.groups_of_edge(edge_select().index().is_in([0, 1])) == {
+            0: ["0"],
+            1: ["0"],
+        }
 
     def test_invalid_groups_of_edge(self) -> None:
         medrecord = create_medrecord()
@@ -1480,25 +1495,37 @@ def test_neighbors(self) -> None:
 
         neighbors = medrecord.neighbors(["0", "1"])
 
-        assert {key: sorted(value) for key, value in neighbors.items()} == {"0": sorted(["1", "3"]), "1": ["0", "2"]}
+        assert {key: sorted(value) for key, value in neighbors.items()} == {
+            "0": sorted(["1", "3"]),
+            "1": ["0", "2"],
+        }
 
         neighbors = medrecord.neighbors(node_select().index().is_in(["0", "1"]))
 
-        assert {key: sorted(value) for key, value in neighbors.items()} == {"0": sorted(["1", "3"]), "1": ["0", "2"]}
+        assert {key: sorted(value) for key, value in neighbors.items()} == {
+            "0": sorted(["1", "3"]),
+            "1": ["0", "2"],
+        }
 
         neighbors = medrecord.neighbors("0", directed=False)
 
         assert sorted(["1", "3"]) == sorted(neighbors)
 
         neighbors = medrecord.neighbors(["0", "1"], directed=False)
 
-        assert {key: sorted(value) for key, value in neighbors.items()} == {"0": sorted(["1", "3"]), "1": ["0", "2"]}
+        assert {key: sorted(value) for key, value in neighbors.items()} == {
+            "0": sorted(["1", "3"]),
+            "1": ["0", "2"],
+        }
 
         neighbors = medrecord.neighbors(
             node_select().index().is_in(["0", "1"]), directed=False
         )
 
-        assert {key: sorted(value) for key, value in neighbors.items()} == {"0": sorted(["1", "3"]), "1": ["0", "2"]}
+        assert {key: sorted(value) for key, value in neighbors.items()} == {
+            "0": sorted(["1", "3"]),
+            "1": ["0", "2"],
+        }
 
     def test_invalid_neighbors(self) -> None:
         medrecord = create_medrecord()