fix: linting issues

medmodels/medrecord/__init__.py

@@ -25,21 +25,21 @@
     "Bool",
     "DateTime",
     "EdgeIndex",
+    "EdgeIndex",
+    "EdgeOperand",
+    "EdgeQuery",
     "Float",
     "GroupSchema",
+    "GroupSchema",
     "Int",
     "MedRecord",
     "NodeIndex",
+    "NodeIndex",
+    "NodeOperand",
+    "NodeQuery",
     "Null",
     "Option",
     "Schema",
-    "GroupSchema",
-    "NodeIndex",
-    "EdgeIndex",
-    "EdgeQuery",
-    "NodeQuery",
-    "NodeOperand",
-    "EdgeOperand",
     "String",
     "Union",
 ]

medmodels/medrecord/_overview.py

@@ -1,21 +1,25 @@
+from __future__ import annotations
+
 import copy
 from datetime import datetime
-from typing import Dict, List, Optional, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Union
 
 import polars as pl
 
 from medmodels.medrecord.schema import AttributesSchema, AttributeType
-from medmodels.medrecord.types import (
-    AttributeInfo,
-    Attributes,
-    EdgeIndex,
-    Group,
-    MedRecordAttribute,
-    NodeIndex,
-    NumericAttributeInfo,
-    StringAttributeInfo,
-    TemporalAttributeInfo,
-)
+
+if TYPE_CHECKING:
+    from medmodels.medrecord.types import (
+        AttributeInfo,
+        Attributes,
+        EdgeIndex,
+        Group,
+        MedRecordAttribute,
+        NodeIndex,
+        NumericAttributeInfo,
+        StringAttributeInfo,
+        TemporalAttributeInfo,
+    )
 
 
 def extract_attribute_summary(
@@ -210,7 +214,7 @@ def prettify_table(
                 row[2] = str(attribute) if first_line else ""
 
                 # displaying info based on the type
-                if "values" in info.keys():
+                if "values" in info:
                     row[3] = info[key]
                 else:
                     if isinstance(info[key], float):

medmodels/medrecord/indexers.py

@@ -2,7 +2,6 @@
 
 from typing import TYPE_CHECKING, Callable, Dict, Tuple, Union, overload
 
-from medmodels.medrecord.querying import EdgeQuery, NodeQuery
 from medmodels.medrecord.types import (
     Attributes,
     AttributesInput,
@@ -22,6 +21,7 @@
 
 if TYPE_CHECKING:
     from medmodels import MedRecord
+    from medmodels.medrecord.querying import EdgeQuery, NodeQuery
 
 
 class NodeIndexer:
@@ -63,7 +63,7 @@ def __getitem__(
         key: Tuple[Union[NodeIndexInputList, NodeQuery, slice], MedRecordAttribute],
     ) -> Dict[NodeIndex, MedRecordValue]: ...
 
-    def __getitem__(
+    def __getitem__(  # noqa: C901
         self,
         key: Union[
             NodeIndex,
@@ -252,7 +252,7 @@ def __setitem__(
         value: MedRecordValue,
     ) -> None: ...
 
-    def __setitem__(
+    def __setitem__(  # noqa: C901
         self,
         key: Union[
             NodeIndex,
@@ -523,7 +523,7 @@ def __setitem__(
             return None
         return None
 
-    def __delitem__(
+    def __delitem__(  # noqa: C901
         self,
         key: Tuple[
             Union[NodeIndex, NodeIndexInputList, NodeQuery, slice],
@@ -712,7 +712,7 @@ def __getitem__(
         key: Tuple[Union[EdgeIndexInputList, EdgeQuery, slice], MedRecordAttribute],
     ) -> Dict[EdgeIndex, MedRecordValue]: ...
 
-    def __getitem__(
+    def __getitem__(  # noqa: C901
         self,
         key: Union[
             EdgeIndex,
@@ -901,7 +901,7 @@ def __setitem__(
         value: MedRecordValue,
     ) -> None: ...
 
-    def __setitem__(
+    def __setitem__(  # noqa: C901
         self,
         key: Union[
             EdgeIndex,
@@ -1170,7 +1170,7 @@ def __setitem__(
             return None
         return None
 
-    def __delitem__(
+    def __delitem__(  # noqa: C901
         self,
         key: Tuple[
             Union[EdgeIndex, EdgeIndexInputList, EdgeQuery, slice],

medmodels/medrecord/medrecord.py

@@ -86,7 +86,7 @@ def __init__(
         data: Dict[Group, AttributeInfo],
         group_header: str,
         decimal: int,
-    ):
+    ) -> None:
         """Initializes the OverviewTable class.
 
         Args:

medmodels/medrecord/querying.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import sys
 from enum import Enum
 from typing import TYPE_CHECKING, Callable, List, Union
 
@@ -27,6 +26,8 @@
 )
 
 if TYPE_CHECKING:
+    import sys
+
     if sys.version_info >= (3, 10):
         from typing import TypeAlias
     else:

medmodels/medrecord/schema.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from enum import Enum, auto
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union, overload
+from typing import TYPE_CHECKING, Dict, Iterator, List, Optional, Tuple, Union, overload
 
 from medmodels._medmodels import (
     PyAttributeDataType,
@@ -143,7 +143,7 @@ def __contains__(self, key: MedRecordAttribute) -> bool:
         """
         return key in self._attributes_schema
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[MedRecordAttribute]:
         """Returns an iterator over the attributes schema.
 
         Returns:
@@ -190,23 +190,23 @@ def __eq__(self, value: object) -> bool:
 
         return True
 
-    def keys(self):
+    def keys(self):  # noqa: ANN201
         """Returns the attribute keys in the schema.
 
         Returns:
             KeysView: A view object displaying a list of dictionary's keys.
         """
         return self._attributes_schema.keys()
 
-    def values(self):
+    def values(self):  # noqa: ANN201
         """Returns the attribute values in the schema.
 
         Returns:
             ValuesView: A view object displaying a list of dictionary's values.
         """
         return self._attributes_schema.values()
 
-    def items(self):
+    def items(self):  # noqa: ANN201
         """Returns the attribute key-value pairs in the schema.
 
         Returns:

medmodels/medrecord/tests/test_builder.py

@@ -1,5 +1,7 @@
 import unittest
 
+import pytest
+
 import medmodels.medrecord as mr
 
 
@@ -60,7 +62,10 @@ def test_with_schema(self) -> None:
 
         medrecord.add_nodes(("node1", {"attribute": 1}))
 
-        with self.assertRaises(ValueError):
+        with pytest.raises(
+            ValueError,
+            match="Attribute attribute of node with index node2 is of type String. Expected Integer.",
+        ):
             medrecord.add_nodes(("node2", {"attribute": "1"}))
 
 

medmodels/medrecord/tests/test_datatype.py

@@ -116,7 +116,7 @@ def test_union(self) -> None:
         assert mr.Union(mr.String(), mr.Int()) != mr.Union(mr.Int(), mr.String())
 
     def test_invalid_union(self) -> None:
-        with pytest.raises(ValueError):
+        with pytest.raises(ValueError, match="Union must have at least two arguments"):
             mr.Union(mr.String())
 
     def test_option(self) -> None: