datasets: fix type errors

l2gv2/datasets/_base.py

@@ -196,7 +196,7 @@ def get_edge_list(
 
     def get_networkx(
         self, temp: bool = True
-    ) -> Union[nx.DiGraph, dict[datetime.datetime, nx.DiGraph]]:
+    ) -> Union[nx.Graph, dict[datetime.datetime, nx.Graph]]:
         """Returns networkx.DiGraph representation
 
         Args:
@@ -205,7 +205,7 @@ def get_networkx(
         """
 
         if self.temporal and temp:
-            nx_graphs = {}
+            nx_graphs: dict[datetime.datetime, nx.Graph] = {}
             for d in tqdm(self.datelist):
                 edges = (
                     self.edges.filter(pl.col("timestamp") == d)
@@ -214,11 +214,10 @@ def get_networkx(
                 )
                 edge_list = [tuple(x) for x in edges]
                 nx_graphs[d] = nx.from_edgelist(edge_list, create_using=nx.DiGraph)
-        else:
-            edges = self.edges.select("source", "dest").unique().to_numpy()
-            edge_list = [tuple(x) for x in edges]
-            nx_graphs = nx.from_edgelist(edge_list, create_using=nx.DiGraph)
-        return nx_graphs
+            return nx_graphs
+        edges = self.edges.select("source", "dest").unique().to_numpy()
+        edge_list = [tuple(x) for x in edges]
+        return nx.from_edgelist(edge_list, create_using=nx.DiGraph)
 
     def get_edge_index(
         self, temp: bool = True

notebooks/asad.ipynb

@@ -353,7 +353,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.0"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,

notebooks/nas.ipynb

@@ -899,7 +899,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.0"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,

notebooks/summary.ipynb

@@ -789,7 +789,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.12.0"
+   "version": "3.11.10"
   }
  },
  "nbformat": 4,