misc: cleanup

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"
  }
  """

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]

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: