Fix mypy complaints

funlib/persistence/arrays/datasets.py

@@ -9,7 +9,7 @@
 import logging
 import os
 import shutil
-from typing import Optional
+from typing import Optional, Union
 
 logger = logging.getLogger(__name__)
 
@@ -180,7 +180,7 @@ def prepare_ds(
     write_roi: Roi = None,
     write_size: Coordinate = None,
     num_channels: Optional[int] = None,
-    compressor: str = "default",
+    compressor: Union[str, dict] = "default",
     delete: bool = False,
     force_exact_write_size: bool = False,
 ) -> Array:

funlib/persistence/graphs/graph_database.py

@@ -33,15 +33,15 @@ def __getitem__(self, roi) -> Graph:
 
     @property
     @abstractmethod
-    def node_attrs(self) -> list[str]:
+    def node_attrs(self) -> dict[str, type]:
         """
         Return the node attributes supported by the database.
         """
         pass
 
     @property
     @abstractmethod
-    def edge_attrs(self) -> list[str]:
+    def edge_attrs(self) -> dict[str, type]:
         """
         Return the edge attributes supported by the database.
         """

funlib/persistence/graphs/pgsql_graph_database.py

@@ -23,7 +23,7 @@ def __init__(
         total_roi: Optional[Roi] = None,
         nodes_table: str = "nodes",
         edges_table: str = "edges",
-        endpoint_names: Optional[tuple[str, str]] = None,
+        endpoint_names: Optional[list[str]] = None,
         node_attrs: Optional[dict[str, type]] = None,
         edge_attrs: Optional[dict[str, type]] = None,
     ):
@@ -122,16 +122,19 @@ def _store_metadata(self, metadata) -> None:
             "metadata", ["value"], [[json.dumps(metadata)]], fail_if_exists=True
         )
 
-    def _read_metadata(self) -> dict[str, Any]:
+    def _read_metadata(self) -> Optional[dict[str, Any]]:
         try:
             self.__exec("SELECT value FROM metadata")
         except psycopg2.errors.UndefinedTable:
             self.connection.rollback()
             return None
 
-        metadata = self.cur.fetchone()[0]
+        result = self.cur.fetchone()
+        if result is not None:
+            metadata = result[0]
+            return json.loads(metadata)
 
-        return json.loads(metadata)
+        return None
 
     def _select_query(self, query) -> Iterable[Any]:
         self.__exec(query)

funlib/persistence/graphs/sql_graph_database.py

@@ -66,7 +66,7 @@ def __init__(
         total_roi: Optional[Roi] = None,
         nodes_table: str = "nodes",
         edges_table: str = "edges",
-        endpoint_names: Optional[tuple[str, str]] = None,
+        endpoint_names: Optional[list[str]] = None,
         node_attrs: Optional[dict[str, type]] = None,
         edge_attrs: Optional[dict[str, type]] = None,
     ):
@@ -101,7 +101,7 @@ def _store_metadata(self, metadata) -> None:
         pass
 
     @abstractmethod
-    def _read_metadata(self) -> dict[str, Any]:
+    def _read_metadata(self) -> Optional[dict[str, Any]]:
         pass
 
     @abstractmethod
@@ -217,7 +217,7 @@ def edge_attrs(self) -> dict[str, type]:
         return self._edge_attrs if self._edge_attrs is not None else {}
 
     @edge_attrs.setter
-    def edge_attrs(self, value: Optional[Iterable[str]]) -> None:
+    def edge_attrs(self, value: dict[str, type]) -> None:
         self._edge_attrs = value
 
     def read_nodes(

funlib/persistence/graphs/sqlite_graph_database.py

@@ -21,7 +21,7 @@ def __init__(
         total_roi: Optional[Roi] = None,
         nodes_table: str = "nodes",
         edges_table: str = "edges",
-        endpoint_names: Optional[tuple[str, str]] = None,
+        endpoint_names: Optional[list[str]] = None,
         node_attrs: Optional[dict[str, type]] = None,
         edge_attrs: Optional[dict[str, type]] = None,
     ):
@@ -87,7 +87,7 @@ def _store_metadata(self, metadata):
         with open(self.meta_collection, "w") as f:
             json.dump(metadata, f)
 
-    def _read_metadata(self) -> dict[str, Any]:
+    def _read_metadata(self) -> Optional[dict[str, Any]]:
         if not self.meta_collection.exists():
             return None
 

pyproject.toml

@@ -31,7 +31,7 @@ dependencies = [
 version = { attr = "funlib.persistence.__version__" }
 
 [project.optional-dependencies]
-dev = ['coverage>=5.0.3', 'pytest', 'black', 'mypy']
+dev = ['coverage>=5.0.3', 'pytest', 'black', 'mypy', 'types-psycopg2']
 
 [tool.black]
 target_version = ['py39', 'py310', 'py311']