From 3498133cfe5044d177cddf4e20f7c2ceec725af2 Mon Sep 17 00:00:00 2001
From: Anthony Mahanna <anthony.mahanna@gmail.com>
Date: Tue, 4 Jul 2023 22:01:38 -0400
Subject: [PATCH] misc: cleanup

---
 adbpyg_adapter/adapter.py | 10 +++++-----
 adbpyg_adapter/typings.py |  4 ++--
 tests/test_adapter.py     |  2 +-
 3 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/adbpyg_adapter/adapter.py b/adbpyg_adapter/adapter.py
index 10d8ada..5a79c31 100644
--- a/adbpyg_adapter/adapter.py
+++ b/adbpyg_adapter/adapter.py
@@ -671,28 +671,28 @@ def ntypes_to_ocollections(
         return list(orphan_collections)
 
     def __fetch_adb_docs(
-        self, col: str, empty_meta: bool, query_options: Any
+        self, col: str, meta_is_empty: bool, query_options: Any
     ) -> DataFrame:
         """Fetches ArangoDB documents within a collection. Returns the
             documents in a DataFrame.
 
         :param col: The ArangoDB collection.
         :type col: str
-        :param empty_meta: Set to True if the metagraph specification
+        :param meta_is_empty: Set to True if the metagraph specification
             for **col** is empty.
-        :type empty_meta: bool
+        :type meta_is_empty: bool
         :param query_options: Keyword arguments to specify AQL query options
             when fetching documents from the ArangoDB instance.
         :type query_options: Any
         :return: A DataFrame representing the ArangoDB documents.
         :rtype: pandas.DataFrame
         """
-        # Only return the entire document if **empty_meta** is False
+        # Only return the entire document if **meta_is_empty** is False
         aql = f"""
             FOR doc IN @@col
                 RETURN {
                     "{ _key: doc._key, _from: doc._from, _to: doc._to }"
-                    if empty_meta
+                    if meta_is_empty
                     else "doc"
                 }
         """
diff --git a/adbpyg_adapter/typings.py b/adbpyg_adapter/typings.py
index 5bb9305..a19cbaa 100644
--- a/adbpyg_adapter/typings.py
+++ b/adbpyg_adapter/typings.py
@@ -8,7 +8,7 @@
     "PyGMap",
 ]
 
-from typing import Any, Callable, DefaultDict, Dict, List, Tuple, Union
+from typing import Any, Callable, DefaultDict, Dict, List, Set, Tuple, Union
 
 from pandas import DataFrame
 from torch import Tensor
@@ -20,7 +20,7 @@
 
 ADBEncoders = Dict[str, DataFrameToTensor]
 ADBMetagraphValues = Union[str, DataFrameToTensor, ADBEncoders]
-ADBMetagraph = Dict[str, Dict[str, Dict[str, ADBMetagraphValues]]]
+ADBMetagraph = Dict[str, Dict[str, Union[Set[str], Dict[str, ADBMetagraphValues]]]]
 
 PyGDataTypes = Union[str, Tuple[str, str, str]]
 PyGMetagraphValues = Union[str, List[str], TensorToDataFrame]
diff --git a/tests/test_adapter.py b/tests/test_adapter.py
index ca5c67e..d1ddf27 100644
--- a/tests/test_adapter.py
+++ b/tests/test_adapter.py
@@ -1066,7 +1066,7 @@ def assert_adb_to_pyg(
 
 
 def assert_adb_to_pyg_meta(
-    meta: Union[str, Dict[str, ADBMetagraphValues]],
+    meta: Union[str, Set[str], Dict[str, ADBMetagraphValues]],
     df: DataFrame,
     pyg_data: Union[NodeStorage, EdgeStorage],
 ) -> None: