chore: import upstream snapshot with attribution

easygraph/datapipe/loader.py

@@ -0,0 +1,90 @@
+import json
+import pickle as pkl
+import re
+
+from pathlib import Path
+from typing import Callable
+from typing import List
+from typing import Optional
+from typing import Union
+
+
+def load_from_pickle(
+    file_path: Path, keys: Optional[Union[str, List[str]]] = None, **kwargs
+):
+    r"""Load data from a pickle file.
+
+    Args:
+        ``file_path`` (``Path``): The local path of the file.
+        ``keys`` (``Union[str, List[str]]``, optional): The keys of the data. Defaults to ``None``.
+    """
+    if isinstance(file_path, list):
+        raise ValueError("This function only support loading data from a single file.")
+    with open(file_path, "rb") as f:
+        data = pkl.load(f, **kwargs)
+    if keys is None:
+        return data
+    elif isinstance(keys, str):
+        return data[keys]
+    else:
+        return {key: data[key] for key in keys}
+
+
+def load_from_json(file_path: Path, **kwargs):
+    r"""Load data from a json file.
+
+    Args:
+        ``file_path`` (``Path``): The local path of the file.
+    """
+    with open(file_path, "r") as f:
+        data = json.load(f, **kwargs)
+    return data
+
+
+def load_from_txt(
+    file_path: Path,
+    dtype: Union[str, Callable],
+    sep: str = ",| |\t",
+    ignore_header: int = 0,
+):
+    r"""Load data from a txt file.
+
+    .. note::
+        The separator is a regular expression of ``re`` module. Multiple separators can be separated by ``|``. More details can refer to `re.split <https://docs.python.org/3/library/re.html#re.split>`_.
+
+    Args:
+        ``file_path`` (``Path``): The local path of the file.
+        ``dtype`` (``Union[str, Callable]``): The data type of the data can be either a string or a callable function.
+        ``sep`` (``str``, optional): The separator of each line in the file. Defaults to ``",| |\t"``.
+        ``ignore_header`` (``int``, optional): The number of lines to ignore in the header of the file. Defaults to ``0``.
+    """
+    cast_fun = ret_cast_fun(dtype)
+    file_path = Path(file_path)
+    assert file_path.exists(), f"{file_path} does not exist."
+    data = []
+    with open(file_path, "r") as f:
+        for _ in range(ignore_header):
+            f.readline()
+        data = [
+            list(map(cast_fun, re.split(sep, line.strip()))) for line in f.readlines()
+        ]
+    return data
+
+
+def ret_cast_fun(dtype: Union[str, Callable]):
+    r"""Return the cast function of the data type. The supported data types are: ``int``, ``float``, ``str``.
+
+    Args:
+        ``dtype`` (``Union[str, Callable]``): The data type of the data can be either a string or a callable function.
+    """
+    if isinstance(dtype, str):
+        if dtype == "int":
+            return int
+        elif dtype == "float":
+            return float
+        elif dtype == "str":
+            return str
+        else:
+            raise ValueError("dtype must be one of 'int', 'float', 'str'.")
+    else:
+        return dtype