quick fix for pandas types on function signatures (#61)

TheJacksonLaboratory · Dec 21, 2022 · f04f8c7 · f04f8c7
1 parent 3bfb442
commit f04f8c7
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 4 deletions.
diff --git a/ezomero/_gets.py b/ezomero/_gets.py
@@ -1114,10 +1114,16 @@ def get_pyramid_levels(conn: BlitzGateway, image_id: int,
     return levels
 
 
+if has_pandas:
+    TableType = pd.core.frame.DataFrame
+else:
+    TableType = List
+
+
 @do_across_groups
 def get_table(conn: BlitzGateway, file_ann_id: int,
               across_groups: Optional[bool] = True
-              ) -> Union[List, pd.core.frame.DataFrame]:
+              ) -> TableType:
     """Get a table from its FileAnnotation object.
 
     Parameters
@@ -1198,7 +1204,7 @@ def get_shape(conn: BlitzGateway, shape_id: int,
 
 
 def _create_table(table_obj: Table
-                  ) -> Union[List, pd.core.frame.DataFrame]:
+                  ) -> TableType:
     if importlib.util.find_spec('pandas'):
         columns = []
         for col in table_obj.getHeaders():

diff --git a/ezomero/_posts.py b/ezomero/_posts.py
@@ -553,7 +553,13 @@ def post_roi(conn: BlitzGateway, image_id: int,
     return roi.getId().getValue()
 
 
-def post_table(conn: BlitzGateway, table: Union[List, pd.core.frame.DataFrame],
+if has_pandas:
+    TableType = pd.core.frame.DataFrame
+else:
+    TableType = List
+
+
+def post_table(conn: BlitzGateway, table: TableType,
                object_type: str, object_id: int,
                title: Optional[str] = "",
                headers: bool = True) -> Union[int, None]:
@@ -637,7 +643,7 @@ def post_table(conn: BlitzGateway, table: Union[List, pd.core.frame.DataFrame],
     return file_ann.id
 
 
-def create_columns(table: Union[List, pd.core.frame.DataFrame],
+def create_columns(table: TableType,
                    headers: bool) -> List[Column]:
     """Helper function to create the correct column types from a table"""
     cols = []