chore: import upstream snapshot with attribution

easygraph/datasets/hypergraph/hypergraph_dataset_base.py

@@ -0,0 +1,119 @@
+from pathlib import Path
+from typing import Any
+from typing import Dict
+from typing import List
+
+from easygraph.datapipe import compose_pipes
+from easygraph.datasets.hypergraph._global import DATASETS_ROOT
+from easygraph.datasets.hypergraph._global import REMOTE_DATASETS_ROOT
+from easygraph.datasets.utils import download_and_check
+
+
+class BaseData:
+    r"""The Base Class of all datasets.
+
+    ::
+
+        self._content = {
+            'item': {
+                'upon': [
+                    {'filename': 'part1.pkl', 'md5': 'xxxxx',},
+                    {'filename': 'part2.pkl', 'md5': 'xxxxx',},
+                ],
+                'loader': loader_function,
+                'preprocess': [datapipe1, datapipe2],
+            },
+            ...
+        }
+
+    """
+
+    def __init__(self, name: str, data_root=None):
+        # configure the data local/remote root
+        self.name = name
+        if data_root is None:
+            self.data_root = DATASETS_ROOT / name
+        else:
+            self.data_root = Path(data_root) / name
+        self.remote_root = REMOTE_DATASETS_ROOT + name + "/"
+        # init
+        self._content = {}
+        self._raw = {}
+
+    def __repr__(self) -> str:
+        return (
+            f"This is {self.name} dataset:\n"
+            + "\n".join(f"  ->  {k}" for k in self.content)
+            + "\nPlease try `data['name']` to get the specified data."
+        )
+
+    @property
+    def content(self):
+        r"""Return the content of the dataset."""
+        return list(self._content.keys())
+
+    def needs_to_load(self, item_name: str) -> bool:
+        r"""Return whether the ``item_name`` of the dataset needs to be loaded.
+
+        Args:
+            ``item_name`` (``str``): The name of the item in the dataset.
+        """
+        assert item_name in self.content, f"{item_name} is not provided in the Data"
+        return (
+            isinstance(self._content[item_name], dict)
+            and "upon" in self._content[item_name]
+            and "loader" in self._content[item_name]
+        )
+
+    def __getitem__(self, key: str) -> Any:
+        if self.needs_to_load(key):
+            cur_cfg = self._content[key]
+            if cur_cfg.get("cache", None) is None:
+                # get raw data
+                item = self.raw(key)
+                # preprocess and cache
+                pipes = cur_cfg.get("preprocess", None)
+                if pipes is not None:
+                    cur_cfg["cache"] = compose_pipes(*pipes)(item)
+                else:
+                    cur_cfg["cache"] = item
+            return cur_cfg["cache"]
+        else:
+            return self._content[key]
+
+    def raw(self, key: str) -> Any:
+        r"""Return the ``key`` of the dataset with un-preprocessed format."""
+        if self.needs_to_load(key):
+            cur_cfg = self._content[key]
+            if self._raw.get(key, None) is None:
+                upon = cur_cfg["upon"]
+                if len(upon) == 0:
+                    return None
+                self.fetch_files(cur_cfg["upon"])
+                file_path_list = [
+                    self.data_root / u["filename"] for u in cur_cfg["upon"]
+                ]
+                if len(file_path_list) == 1:
+                    self._raw[key] = cur_cfg["loader"](file_path_list[0])
+                else:
+                    # here, you should implement a multi-file loader
+                    self._raw[key] = cur_cfg["loader"](file_path_list)
+            return self._raw[key]
+        else:
+            return self._content[key]
+
+    def fetch_files(self, files: List[Dict[str, str]]):
+        r"""Download and check the files if they are not exist.
+
+        Args:
+            ``files`` (``List[Dict[str, str]]``): The files to download, each element
+                in the list is a dict with at lease two keys: ``filename`` and ``md5``.
+                If extra key ``bk_url`` is provided, it will be used to download the
+                file from the backup url.
+        """
+        for file in files:
+            cur_filename = file["filename"]
+            cur_url = file.get("bk_url", None)
+            if cur_url is None:
+                cur_url = self.remote_root + cur_filename
+            download_and_check(cur_url, self.data_root / cur_filename, file["md5"])