Repository: a-r-j/graphein (language: python)

Issue:
NameError for missing `os` import in `graphein.protein.graphs` 
When using the `construct_graph` function in a jupyter notebook, I encountered this error: 

```
File ~/env/python3.10/site-packages/graphein/protein/graphs.py:764, in construct_graph(config, name, pdb_path, uniprot_id, pdb_code, df, chain_selection, model_index, df_processing_funcs, edge_construction_funcs, edge_annotation_funcs, node_annotation_funcs, graph_annotation_funcs, verbose)
    761 if pdb_path is not None and isinstance(pdb_path, Path):
    762     pdb_path = os.fsdecode(pdb_path)
...
--> 764 # If no config is provided, use default
    765 if config is None:
    766     config = ProteinGraphConfig()

NameError: name 'os' is not defined
```

In the `graphs.py` file, there is no `import os`.  Adding this fixes the error.  Not sure what's going on as I've never encountered this before? 




Discussion / hints:

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 1f3387e9..d4a6d71c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -9,7 +9,7 @@ repos:
     -   id: check-yaml
     -   id: check-added-large-files
 -   repo: https://github.com/psf/black
-    rev: 23.1.0
+    rev: 23.3.0
     hooks:
     - id: black
 -   repo: https://github.com/pre-commit/pre-commit-hooks
diff --git a/.requirements/base.in b/.requirements/base.in
index 5f29ea18..ca1ada0d 100644
--- a/.requirements/base.in
+++ b/.requirements/base.in
@@ -1,4 +1,5 @@
-biopandas>=0.4.1
+pandas<2.0.0
+biopandas==0.5.0.dev0
 biopython
 bioservices>=1.10.0
 deepdiff
diff --git a/.requirements/torch_cpu.in b/.requirements/torch_cpu.in
new file mode 100644
index 00000000..d27fa26c
--- /dev/null
+++ b/.requirements/torch_cpu.in
@@ -0,0 +1,1 @@
+einops
diff --git a/.requirements/torch_gpu.in b/.requirements/torch_gpu.in
new file mode 100644
index 00000000..d27fa26c
--- /dev/null
+++ b/.requirements/torch_gpu.in
@@ -0,0 +1,1 @@
+einops
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 86b123b1..eaad192e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,10 +1,5 @@
 ### 1.6.1 - UNRELEASED
 
-* `Protein` tensors have coordinates renamed from `Protein.x` to `Protein.coords`. [#272](https://github.com/a-r-j/graphein/pull/272)
-* Tensor types are now defined using [`jaxtyping`](https://github.com/google/jaxtyping), removing the `torchtyping` dependency [#272](https://github.com/a-r-j/graphein/pull/272)
-* Drops explicit Python 3.7 support. Colab now runs on 3.8+. [#272](https://github.com/a-r-j/graphein/pull/272)
-* Dockerfile now builds from `pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime` (replaces `pytorch/pytorch:1.9.1-cuda11.1-cudnn8-runtime`) [#272](https://github.com/a-r-j/graphein/pull/272)
-
 #### New Features
 
 * [PDBManager] - [#272](https://github.com/a-r-j/graphein/pull/272) Adds a utility for creating custom dataset splits from the PDB.
@@ -12,6 +7,18 @@
 * [ESM] - [#284](https://github.com/a-r-j/graphein/pull/284) - Wrapper for ESMFold batch folding & embedding.
 * [Downloads] MMTF downloading now supported in download utilities. [#272](https://github.com/a-r-j/graphein/pull/272)
 
+#### API Changes
+* The `pdb_path` argument to many functions (e.g. `graphein.protein.graphs.construct_graph`) has been renamed to `path` as this can now accept MMTF files in addition to PDB files.
+* `Protein` tensors have coordinates renamed from `Protein.x` to `Protein.coords`. [#272](https://github.com/a-r-j/graphein/pull/272)
+
+#### Other changes
+* Tensor types are now defined using [`jaxtyping`](https://github.com/google/jaxtyping), removing the `torchtyping` dependency [#272](https://github.com/a-r-j/graphein/pull/272)
+* Drops explicit Python 3.7 support. Colab now runs on 3.8+. [#272](https://github.com/a-r-j/graphein/pull/272)
+* Dockerfile now builds from `pytorch/pytorch:1.13.0-cuda11.6-cudnn8-runtime` (replaces `pytorch/pytorch:1.9.1-cuda11.1-cudnn8-runtime`) [#272](https://github.com/a-r-j/graphein/pull/272)
+* Missing `os` import fixed in [#297(https://github.com/a-r-j/graphein/pull/297). Fixes [#296](https://github.com/a-r-j/graphein/issues/296)
+
+
+
 ### 1.6.0 - 18/03/2023
 
 #### New Features
diff --git a/README.md b/README.md
index 65f75807..07ee0897 100644
--- a/README.md
+++ b/README.md
@@ -89,7 +89,7 @@ from graphein.protein.utils import download_alphafold_structure
 
 config = ProteinGraphConfig()
 fp = download_alphafold_structure("Q5VSL9", aligned_score=False)
-g = construct_graph(config=config, pdb_path=fp)
+g = construct_graph(config=config, path=fp)
 ```
 
 ### Creating a Protein Mesh
diff --git a/graphein/cli.py b/graphein/cli.py
index 0d7689c4..b02784be 100644
--- a/graphein/cli.py
+++ b/graphein/cli.py
@@ -21,7 +21,7 @@
 )
 @click.option(
     "-p",
-    "--pdb_path",
+    "--path",
     help="Path to input pdbs",
     type=click.Path(
         exists=True, file_okay=True, dir_okay=True, path_type=pathlib.Path
@@ -35,20 +35,20 @@
         exists=True, file_okay=False, dir_okay=True, path_type=pathlib.Path
     ),
 )
-def main(config_path, pdb_path, output_path):
+def main(config_path, path, output_path):
     """Build the graphs and save them in output dir."""
     config = parse_config(path=config_path) if config_path else None
-    if pdb_path.is_file():
-        pdb_paths = [pdb_path]
-    elif pdb_path.is_dir():
-        pdb_paths = list(pdb_path.glob("*.pdb"))
+    if path.is_file():
+        paths = [path]
+    elif path.is_dir():
+        paths = list(path.glob("*.pdb"))
     else:
         raise NotImplementedError(
             "Given PDB path needs to point to either a pdb file or a directory with pdb files."
         )
 
-    for path in pdb_paths:
-        g = construct_graph(config=config, pdb_path=str(path))
+    for path in paths:
+        g = construct_graph(config=config, path=str(path))
 
         with open(str(output_path / f"{path.stem}.gpickle"), "wb") as f:
             pickle.dump(g, f)
diff --git a/graphein/ml/datasets/torch_geometric_dataset.py b/graphein/ml/datasets/torch_geometric_dataset.py
index c97196e7..e9e8b030 100644
--- a/graphein/ml/datasets/torch_geometric_dataset.py
+++ b/graphein/ml/datasets/torch_geometric_dataset.py
@@ -41,7 +41,7 @@ def __init__(
         self,
         root: str,
         name: str,
-        pdb_paths: Optional[List[str]] = None,
+        paths: Optional[List[str]] = None,
         pdb_codes: Optional[List[str]] = None,
         uniprot_ids: Optional[List[str]] = None,
         graph_label_map: Optional[Dict[str, torch.Tensor]] = None,
@@ -73,9 +73,9 @@ def __init__(
         :type root: str
         :param name: Name of the dataset. Will be saved to ``data_$name.pt``.
         :type name: str
-        :param pdb_paths: List of full path of pdb files to load. Defaults to
-            ``None``.
-        :type pdb_paths: Optional[List[str]], optional
+        :param paths: List of full path of PDB or MMTF files to load. Defaults
+            to ``None``.
+        :type paths: Optional[List[str]], optional
         :param pdb_codes: List of PDB codes to download and parse from the PDB.
             Defaults to None.
         :type pdb_codes: Optional[List[str]], optional
@@ -139,8 +139,8 @@ def __init__(
             else None
         )
 
-        self.pdb_paths = pdb_paths
-        if self.pdb_paths is None:
+        self.paths = paths
+        if self.paths is None:
             if self.pdb_codes and self.uniprot_ids:
                 self.structures = self.pdb_codes + self.uniprot_ids
             elif self.pdb_codes:
@@ -150,12 +150,12 @@ def __init__(
         # Use local saved pdb_files instead of download or move them to
         # self.root/raw dir
         else:
-            if isinstance(self.pdb_paths, list):
+            if isinstance(self.paths, list):
                 self.structures = [
-                    os.path.splitext(os.path.split(pdb_path)[-1])[0]
-                    for pdb_path in self.pdb_paths
+                    os.path.splitext(os.path.split(path)[-1])[0]
+                    for path in self.paths
                 ]
-                self.pdb_path, _ = os.path.split(self.pdb_paths[0])
+                self.path, _ = os.path.split(self.paths[0])
 
         if self.pdb_codes and self.uniprot_ids:
             self.structures = self.pdb_codes + self.uniprot_ids
@@ -201,8 +201,8 @@ def processed_file_names(self) -> List[str]:
 
     @property
     def raw_dir(self) -> str:
-        if self.pdb_paths is not None:
-            return self.pdb_path  # replace raw dir with user local pdb_path
+        if self.paths is not None:
+            return self.path  # replace raw dir with user local path
         else:
             return os.path.join(self.root, "raw")
 
@@ -276,7 +276,7 @@ def process(self):
         # Create graph objects
         print("Constructing Graphs...")
         graphs = construct_graphs_mp(
-            pdb_path_it=structure_files,
+            path_it=structure_files,
             config=self.config,
             chain_selections=chain_selections,
             return_dict=True,
@@ -329,7 +329,7 @@ class ProteinGraphDataset(Dataset):
     def __init__(
         self,
         root: str,
-        pdb_paths: Optional[List[str]] = None,
+        paths: Optional[List[str]] = None,
         pdb_codes: Optional[List[str]] = None,
         uniprot_ids: Optional[List[str]] = None,
         # graph_label_map: Optional[Dict[str, int]] = None,
@@ -358,9 +358,9 @@ def __init__(
 
         :param root: Root directory where the dataset should be saved.
         :type root: str
-        :param pdb_paths: List of full path of pdb files to load. Defaults to
-            ``None``.
-        :type pdb_paths: Optional[List[str]], optional
+        :param paths: List of full path of PDB or MMTF files to load. Defaults
+            to ``None``.
+        :type paths: Optional[List[str]], optional
         :param pdb_codes: List of PDB codes to download and parse from the PDB.
             Defaults to ``None``.
         :type pdb_codes: Optional[List[str]], optional
@@ -422,8 +422,8 @@ def __init__(
             if uniprot_ids is not None
             else None
         )
-        self.pdb_paths = pdb_paths
-        if self.pdb_paths is None:
+        self.paths = paths
+        if self.paths is None:
             if self.pdb_codes and self.uniprot_ids:
                 self.structures = self.pdb_codes + self.uniprot_ids
             elif self.pdb_codes:
@@ -433,12 +433,12 @@ def __init__(
         # Use local saved pdb_files instead of download or move them to
         # self.root/raw dir
         else:
-            if isinstance(self.pdb_paths, list):
+            if isinstance(self.paths, list):
                 self.structures = [
-                    os.path.splitext(os.path.split(pdb_path)[-1])[0]
-                    for pdb_path in self.pdb_paths
+                    os.path.splitext(os.path.split(path)[-1])[0]
+                    for path in self.paths
                 ]
-                self.pdb_path, _ = os.path.split(self.pdb_paths[0])
+                self.path, _ = os.path.split(self.paths[0])
 
         # Labels & Chains
 
@@ -496,8 +496,8 @@ def processed_file_names(self) -> List[str]:
 
     @property
     def raw_dir(self) -> str:
-        if self.pdb_paths is not None:
-            return self.pdb_path  # replace raw dir with user local pdb_path
+        if self.paths is not None:
+            return self.path  # replace raw dir with user local path
         else:
             return os.path.join(self.root, "raw")
 
@@ -607,7 +607,7 @@ def divide_chunks(l: List[str], n: int = 2) -> Generator:
             file_names = [f"{self.raw_dir}/{pdb}.pdb" for pdb in pdbs]
 
             graphs = construct_graphs_mp(
-                pdb_path_it=file_names,
+                path_it=file_names,
                 config=self.config,
                 chain_selections=chain_selections,
                 return_dict=False,
diff --git a/graphein/protein/features/nodes/dssp.py b/graphein/protein/features/nodes/dssp.py
index 1b28cf0d..d68615b2 100644
--- a/graphein/protein/features/nodes/dssp.py
+++ b/graphein/protein/features/nodes/dssp.py
@@ -85,7 +85,7 @@ def add_dssp_df(
 
     config = G.graph["config"]
     pdb_code = G.graph["pdb_code"]
-    pdb_path = G.graph["pdb_path"]
+    path = G.graph["path"]
     pdb_name = G.graph["name"]
 
     # Extract DSSP executable
@@ -97,9 +97,9 @@ def add_dssp_df(
     ), "DSSP must be on PATH and marked as an executable"
 
     pdb_file = None
-    if pdb_path:
-        if os.path.isfile(pdb_path):
-            pdb_file = pdb_path
+    if path:
+        if os.path.isfile(path):
+            pdb_file = path
     else:
         if config.pdb_dir:
             if os.path.isfile(config.pdb_dir / (pdb_code + ".pdb")):
diff --git a/graphein/protein/graphs.py b/graphein/protein/graphs.py
index b780b34d..bccff6ef 100644
--- a/graphein/protein/graphs.py
+++ b/graphein/protein/graphs.py
@@ -7,6 +7,7 @@
 # Code Repository: https://github.com/a-r-j/graphein
 from __future__ import annotations
 
+import os
 import traceback
 from contextlib import nullcontext
 from functools import partial
@@ -17,6 +18,7 @@
 import numpy as np
 import pandas as pd
 from Bio.PDB.Polypeptide import three_to_one
+from biopandas.mmtf import PandasMmtf
 from biopandas.pdb import PandasPdb
 from loguru import logger as log
 from rich.progress import Progress
@@ -68,7 +70,7 @@ def subset_structure_to_rna(
 
 
 def read_pdb_to_dataframe(
-    pdb_path: Optional[os.Pathlike] = None,
+    path: Optional[Union[str, os.PathLike]] = None,
     pdb_code: Optional[str] = None,
     uniprot_id: Optional[str] = None,
     model_index: int = 1,
@@ -80,8 +82,8 @@ def read_pdb_to_dataframe(
     their cartesian coordinates in 3D space. Also contains associated metadata
     from the PDB file.
 
-    :param pdb_path: path to PDB file. Defaults to ``None``.
-    :type pdb_path: str, optional
+    :param path: path to PDB or MMTF file. Defaults to ``None``.
+    :type path: str, optional
     :param pdb_code: 4-character PDB accession. Defaults to ``None``.
     :type pdb_code: str, optional
     :param uniprot_id: UniProt ID to build graph from AlphaFoldDB. Defaults to
@@ -93,18 +95,23 @@ def read_pdb_to_dataframe(
     :returns: ``pd.DataFrame`` containing protein structure
     :rtype: pd.DataFrame
     """
-    if pdb_code is None and pdb_path is None and uniprot_id is None:
+    if pdb_code is None and path is None and uniprot_id is None:
         raise NameError(
-            "One of pdb_code, pdb_path or uniprot_id must be specified!"
+            "One of pdb_code, path or uniprot_id must be specified!"
         )
 
-    if pdb_path is not None:
-        if isinstance(pdb_path, Path):
-            pdb_path = os.fsdecode(pdb_path)
-        atomic_df = PandasPdb().read_pdb(pdb_path)
+    if path is not None:
+        if isinstance(path, Path):
+            path = os.fsdecode(path)
+        if path.endswith(".pdb"):
+            atomic_df = PandasPdb().read_pdb(path)
+        elif path.endswith(".mmtf"):
+            atomic_df = PandasMmtf().read_mmtf(path)
+        else:
+            raise ValueError("File must be either .pdb or .mmtf")
     elif uniprot_id is not None:
         atomic_df = PandasPdb().fetch_pdb(
-            uniprot_id=uniprot_id, source="alphafold2-v2"
+            uniprot_id=uniprot_id, source="alphafold2-v3"
         )
     else:
         atomic_df = PandasPdb().fetch_pdb(pdb_code)
@@ -486,7 +493,7 @@ def initialise_graph_with_metadata(
     granularity: str,
     name: Optional[str] = None,
     pdb_code: Optional[str] = None,
-    pdb_path: Optional[str] = None,
+    path: Optional[str] = None,
 ) -> nx.Graph:
     """
     Initializes the nx Graph object with initial metadata.
@@ -507,26 +514,26 @@ def initialise_graph_with_metadata(
     :param pdb_code: PDB ID / Accession code, if the PDB is available on the
         PDB database.
     :type pdb_code: Optional[str], defaults to ``None``.
-    :param pdb_path: path to local PDB file, if constructing a graph from a
+    :param path: path to local PDB or MMTF file, if constructing a graph from a
         local file.
-    :type pdb_path: Optional[str], defaults to ``None``.
+    :type path: Optional[str], defaults to ``None``.
     :return: Returns initial protein structure graph with metadata.
     :rtype: nx.Graph
     """
-    if pdb_path is not None and isinstance(pdb_path, Path):
-        pdb_path = os.fsdecode(pdb_path)
+    if path is not None and isinstance(path, Path):
+        path = os.fsdecode(path)
 
     # Get name for graph if no name was provided
     if name is None:
-        if pdb_path is not None:
-            name = get_protein_name_from_filename(pdb_path)
+        if path is not None:
+            name = get_protein_name_from_filename(path)
         else:
             name = pdb_code
 
     G = nx.Graph(
         name=name,
         pdb_code=pdb_code,
-        pdb_path=pdb_path,
+        path=path,
         chain_ids=list(protein_df["chain_id"].unique()),
         pdb_df=protein_df,
         raw_pdb_df=raw_pdb_df,
@@ -679,7 +686,7 @@ def compute_edges(
 def construct_graph(
     config: Optional[ProteinGraphConfig] = None,
     name: Optional[str] = None,
-    pdb_path: Optional[os.Pathlike] = None,
+    path: Optional[Union[str, os.PathLike]] = None,
     uniprot_id: Optional[str] = None,
     pdb_code: Optional[str] = None,
     df: Optional[pd.DataFrame] = None,
@@ -693,7 +700,7 @@ def construct_graph(
     verbose: bool = True,
 ) -> nx.Graph:
     """
-    Constructs protein structure graph from a ``pdb_code``, ``pdb_path``,
+    Constructs protein structure graph from a ``pdb_code``, ``path``,
     ``uniprot_id`` or a BioPandas DataFrame containing ``ATOM`` data.
 
     Users can provide a :class:`~graphein.protein.config.ProteinGraphConfig`
@@ -708,9 +715,9 @@ def construct_graph(
     :param name: an optional given name for the graph. the PDB ID or PDB file
         name will be used if not specified.
     :type name: str, optional
-    :param pdb_path: Path to ``pdb_file`` when constructing a graph from a
+    :param path: Path to PDB or MMTF file when constructing a graph from a
         local pdb file. Default is ``None``.
-    :type pdb_path: Optional[str], defaults to ``None``
+    :type path: Optional[str], defaults to ``None``
     :param pdb_code: A 4-character PDB ID / accession to be used to construct
         the graph, if available. Default is ``None``.
     :type pdb_code: Optional[str], defaults to ``None``
@@ -748,18 +755,13 @@ def construct_graph(
     :rtype: nx.Graph
     """
 
-    if (
-        pdb_code is None
-        and pdb_path is None
-        and uniprot_id is None
-        and df is None
-    ):
+    if pdb_code is None and path is None and uniprot_id is None and df is None:
         raise ValueError(
             "Either a PDB ID, UniProt ID, a dataframe or a path to a local PDB file"
             " must be specified to construct a graph"
         )
-    if pdb_path is not None and isinstance(pdb_path, Path):
-        pdb_path = os.fsdecode(pdb_path)
+    if path is not None and isinstance(path, Path):
+        path = os.fsdecode(path)
 
     # If no config is provided, use default
     if config is None:
@@ -800,7 +802,7 @@ def construct_graph(
         )
         if df is None:
             raw_df = read_pdb_to_dataframe(
-                pdb_path,
+                path,
                 pdb_code,
                 uniprot_id,
                 model_index=model_index,
@@ -832,7 +834,7 @@ def construct_graph(
             raw_pdb_df=raw_df,
             name=name,
             pdb_code=pdb_code,
-            pdb_path=pdb_path,
+            path=path,
             granularity=config.granularity,
         )
         # Add nodes to graph
@@ -878,8 +880,8 @@ def _mp_graph_constructor(
 
     :param args: Tuple of pdb code/path and the chain selection for that PDB.
     :type args: Tuple[str, str]
-    :param use_pdb_code: Whether we are using ``"pdb_code"``s, ``pdb_path``s or
-        ``"uniprot_id"``s.
+    :param use_pdb_code: Whether we are using ``"pdb_code"``s, ``path``s
+        (to PDB or MMTF files) or ``"uniprot_id"``s.
     :type use_pdb_code: bool
     :param config: Protein structure graph construction config
         (see: :class:`graphein.protein.config.ProteinGraphConfig`).
@@ -897,9 +899,9 @@ def _mp_graph_constructor(
             return func(
                 pdb_code=args[0], chain_selection=args[1], model_index=args[2]
             )
-        elif source == "pdb_path":
+        elif source == "path":
             return func(
-                pdb_path=args[0], chain_selection=args[1], model_index=args[2]
+                path=args[0], chain_selection=args[1], model_index=args[2]
             )
         elif source == "uniprot_id":
             return func(
@@ -919,7 +921,7 @@ def _mp_graph_constructor(
 
 def construct_graphs_mp(
     pdb_code_it: Optional[List[str]] = None,
-    pdb_path_it: Optional[List[str]] = None,
+    path_it: Optional[List[str]] = None,
     uniprot_id_it: Optional[List[str]] = None,
     chain_selections: Optional[List[str]] = None,
     model_indices: Optional[List[str]] = None,
@@ -934,9 +936,9 @@ def construct_graphs_mp(
 
     :param pdb_code_it: List of pdb codes to use for protein graph construction
     :type pdb_code_it: Optional[List[str]], defaults to ``None``
-    :param pdb_path_it: List of paths to PDB files to use for protein graph
+    :param path_it: List of paths to PDB or MMTF files to use for protein graph
         construction.
-    :type pdb_path_it: Optional[List[str]], defaults to ``None``
+    :type path_it: Optional[List[str]], defaults to ``None``
     :param chain_selections: List of chains to select from the protein
         structures (e.g. ``["ABC", "A", "L", "CD"...]``).
     :type chain_selections: Optional[List[str]], defaults to ``None``
@@ -961,16 +963,16 @@ def construct_graphs_mp(
     :rtype: Union[List[nx.Graph], Dict[str, nx.Graph]]
     """
     assert (
-        pdb_code_it is not None or pdb_path_it is not None
+        pdb_code_it is not None or path_it is not None
     ), "Iterable of pdb codes, pdb paths or uniprot IDs required."
 
     if pdb_code_it is not None:
         pdbs = pdb_code_it
         source = "pdb_code"
 
-    if pdb_path_it is not None:
-        pdbs = pdb_path_it
-        source = "pdb_path"
+    if path_it is not None:
+        pdbs = path_it
+        source = "path"
 
     if uniprot_id_it is not None:
         pdbs = uniprot_id_it
diff --git a/graphein/protein/tensor/data.py b/graphein/protein/tensor/data.py
index da8d786b..3addd7d1 100644
--- a/graphein/protein/tensor/data.py
+++ b/graphein/protein/tensor/data.py
@@ -101,7 +101,7 @@ class Protein(Data):
 
         # From a PyG Data object
         data = gpt.io.protein_to_pyg(
-            pdb_code="4hhb", # Can alternatively pass a path or a uniprot ID (for AF2) with pdb_path=... and uniprot_id=...
+            pdb_code="4hhb", # Can alternatively pass a path or a uniprot ID (for AF2) with path=... and uniprot_id=...
             chain_selection="ABCD", # Select all 4 chains
             deprotonate=True, # Deprotonate the structure
             keep_insertions=False, # Remove insertions
@@ -225,7 +225,7 @@ def from_data(self, data: Data) -> "Protein":
             >>> import graphein.protein.tensor as gpt
 
             >>> data = gpt.io.protein_to_pyg(
-                ... pdb_code="4hhb", # Can alternatively pass a path or a uniprot ID (for AF2) with pdb_path=... and uniprot_id=...
+                ... pdb_code="4hhb", # Can alternatively pass a path or a uniprot ID (for AF2) with path=... and uniprot_id=...
                 ... chain_selection="ABCD", # Select all 4 chains
                 ... deprotonate=True, # Deprotonate the structure
                 ... keep_insertions=False, # Remove insertions
@@ -376,7 +376,7 @@ def from_pdb_code(
 
     def from_pdb_file(
         self,
-        pdb_path: str,
+        path: str,
         chain_selection: str = "all",
         deprotonate: bool = True,
         keep_insertions=False,
@@ -387,7 +387,7 @@ def from_pdb_file(
         graph_labels: Optional[torch.Tensor] = None,
     ) -> "Protein":
         data = protein_to_pyg(
-            pdb_path=pdb_path,
+            path=path,
             chain_selection=chain_selection,
             deprotonate=deprotonate,
             keep_insertions=keep_insertions,
@@ -887,7 +887,7 @@ def from_pdb_codes(
 
     def from_pdb_files(
         self,
-        pdb_paths: List[str],
+        paths: List[str],
         chain_selection: Optional[List[str]] = None,
         node_labels: Optional[List[torch.Tensor]] = None,
         graph_labels: Optional[List[torch.Tensor]] = None,
@@ -899,7 +899,7 @@ def from_pdb_files(
     ):
         proteins = [
             Protein().from_pdb_file(
-                pdb_path=pdb,
+                path=pdb,
                 chain_selection=chain_selection[i]
                 if chain_selection is not None
                 else "all",
@@ -915,7 +915,7 @@ def from_pdb_files(
                 if graph_labels is not None
                 else None,
             )
-            for i, pdb in enumerate(pdb_paths)
+            for i, pdb in enumerate(paths)
         ]
         batch = Batch.from_data_list(proteins)
         self.from_batch(batch)
@@ -1297,7 +1297,7 @@ def apply_structural_noise(
 
 
 def to_protein(
-    pdb_path: Optional[str] = None,
+    path: Optional[str] = None,
     pdb_code: Optional[str] = None,
     uniprot_id: Optional[str] = None,
     chain_selection: str = "all",
@@ -1321,7 +1321,7 @@ def to_protein(
         gpt.data.to_protein(pdb_code="3eiy", ...)
 
         # From PDB Path
-        gpt.io.to_protein(pdb_path="3eiy.pdb", ...)
+        gpt.io.to_protein(path="3eiy.pdb", ...)
 
         # From UniProt ID
         gpt.io.to_protein(uniprot_id="Q5VSL9", ...)
@@ -1332,7 +1332,7 @@ def to_protein(
         :func:`graphein.protein.tensor.data.to_protein_mp`
 
 
-    :param pdb_path: Path to PDB file. Default is ``None``.
+    :param path: Path to PDB or MMTF file. Default is ``None``.
     :param pdb_code: PDB accesion code. Default is ``None``.
     :param uniprot_id: UniProt ID. Default is ``None``.
     :param chain_selection: Selection of chains to include (e.g. ``"ABC"``) or
@@ -1349,7 +1349,7 @@ def to_protein(
     :rtype: Protein
     """
     data = protein_to_pyg(
-        pdb_path=pdb_path,
+        path=path,
         pdb_code=pdb_code,
         uniprot_id=uniprot_id,
         chain_selection=chain_selection,
@@ -1385,9 +1385,9 @@ def _mp_constructor(
                 pdb_code=args[0],
                 chain_selection=args[1],  # , model_index=args[2]
             )
-        elif source == "pdb_path":
+        elif source == "path":
             return func(
-                pdb_path=args[0],
+                path=args[0],
                 chain_selection=args[1],  # , model_index=args[2]
             )
         elif source == "uniprot_id":
@@ -1405,7 +1405,7 @@ def _mp_constructor(
 
 
 def to_protein_mp(
-    pdb_paths: Optional[str] = None,
+    paths: Optional[str] = None,
     pdb_codes: Optional[str] = None,
     uniprot_ids: Optional[str] = None,
     chain_selections: Optional[List[str]] = None,
@@ -1430,7 +1430,7 @@ def to_protein_mp(
         gpt.data.to_protein_mp(pdb_codes=["3eiy", "4hhb", ..., num_cores=8])
 
         # From PDB Paths
-        gpt.io.to_protein_mp(pdb_paths=["3eiy.pdb", "4hhb.pdb", ...])
+        gpt.io.to_protein_mp(paths=["3eiy.pdb", "4hhb.pdb", ...])
 
         # From UniProt IDs
         gpt.io.to_protein_mp(uniprot_ids=["Q5VSL9", ...])
@@ -1442,7 +1442,7 @@ def to_protein_mp(
         :func:`graphein.protein.tensor.data.to_protein`
 
 
-    :param pdb_paths: Path to PDB file. Default is ``None``.
+    :param paths: Path to PDB or MMTF files. Default is ``None``.
     :param pdb_codes: PDB accesion code. Default is ``None``.
     :param uniprot_ids: UniProt ID. Default is ``None``.
     :param chain_selections: Selection of chains to include (e.g. ``"ABC"``) or
@@ -1460,18 +1460,16 @@ def to_protein_mp(
     :rtype: List[Protein]
     """
     assert (
-        pdb_codes is not None
-        or pdb_paths is not None
-        or uniprot_ids is not None
+        pdb_codes is not None or paths is not None or uniprot_ids is not None
     ), "Iterable of pdb codes, pdb paths or uniprot IDs required."
 
     if pdb_codes is not None:
         pdbs = pdb_codes
         source = "pdb_code"
 
-    if pdb_paths is not None:
-        pdbs = pdb_paths
-        source = "pdb_path"
+    if paths is not None:
+        pdbs = paths
+        source = "path"
 
     if uniprot_ids is not None:
         pdbs = uniprot_ids
diff --git a/graphein/protein/tensor/dataset.py b/graphein/protein/tensor/dataset.py
index 83381128..e32bd486 100644
--- a/graphein/protein/tensor/dataset.py
+++ b/graphein/protein/tensor/dataset.py
@@ -48,7 +48,7 @@ def __init__(
         pdb_dir: str,
         out_dir: str,
         overwrite: bool = False,
-        pdb_paths: Optional[List[str]] = None,
+        paths: Optional[List[str]] = None,
         pdb_codes: Optional[List[str]] = None,
         uniprot_ids: Optional[List[str]] = None,
         sequences: Optional[Union[List[str], str]] = None,
@@ -81,9 +81,9 @@ def __init__(
         :param overwrite: Whether to overwrite existing files. Defaults to
             ``False``.
         :type overwrite: bool, optional
-        :param pdb_paths: List of full path of pdb files to load. Defaults to
+        :param paths: List of full path of PDB or MMTF files to load. Defaults to
             ``None``.
-        :type pdb_paths: Optional[List[str]], optional
+        :type paths: Optional[List[str]], optional
         :param pdb_codes: List of PDB codes to download and parse from the PDB.
             Defaults to ``None``.
         :type pdb_codes: Optional[List[str]], optional
@@ -153,8 +153,8 @@ def __init__(
             self.sequences = sequences
             self.fold_sequences()
 
-        self.pdb_paths = pdb_paths
-        if self.pdb_paths is None:
+        self.paths = paths
+        if self.paths is None:
             if self.pdb_codes and self.uniprot_ids:
                 self.structures = self.pdb_codes + self.uniprot_ids
             elif self.pdb_codes:
@@ -164,12 +164,12 @@ def __init__(
         # Use local saved pdb_files instead of download or move them to
         # self.root/raw dir
         else:
-            if isinstance(self.pdb_paths, list):
+            if isinstance(self.paths, list):
                 self.structures = [
-                    os.path.splitext(os.path.split(pdb_path)[-1])[0]
-                    for pdb_path in self.pdb_paths
+                    os.path.splitext(os.path.split(path)[-1])[0]
+                    for path in self.paths
                 ]
-                self.pdb_path, _ = os.path.split(self.pdb_paths[0])
+                self.path, _ = os.path.split(self.paths[0])
 
         # Labels & chains
         self.chain_selections = chain_selections
@@ -410,7 +410,7 @@ def divide_chunks(l: List[str], n: int = 2) -> Generator:
             ]
             print(file_names)
             graphs = to_protein_mp(
-                pdb_paths=file_names,
+                paths=file_names,
                 # config=self.config,
                 # TODO args
                 chain_selections=chains,
diff --git a/graphein/protein/tensor/io.py b/graphein/protein/tensor/io.py
index aaa34305..d5761c78 100644
--- a/graphein/protein/tensor/io.py
+++ b/graphein/protein/tensor/io.py
@@ -5,6 +5,7 @@
 # Project Website: https://github.com/a-r-j/graphein
 # Code Repository: https://github.com/a-r-j/graphein
 
+import os
 from typing import List, Optional, Union
 
 import numpy as np
@@ -94,7 +95,7 @@ def get_protein_length(df: pd.DataFrame, insertions: bool = False) -> int:
 
 
 def protein_to_pyg(
-    pdb_path: Optional[str] = None,
+    path: Optional[Union[str, os.PathLike]] = None,
     pdb_code: Optional[str] = None,
     uniprot_id: Optional[str] = None,
     chain_selection: str = "all",
@@ -117,13 +118,16 @@ def protein_to_pyg(
         gpt.io.protein_to_pyg(pdb_code="3eiy")
 
         # From PDB Path
-        gpt.io.protein_to_pyg(pdb_path="3eiy.pdb")
+        gpt.io.protein_to_pyg(path="3eiy.pdb")
+
+        # From MMTF Path
+        gpt.io.protein_to_pyg(path="3eiy.mmtf")
 
         # From UniProt ID
-        gpt.io.protein_to_pyg(pdb_path="Q5VSL9")
+        gpt.io.protein_to_pyg(uniprot_id="Q5VSL9")
 
 
-    :param pdb_path: Path to PDB file. Default is ``None``.
+    :param path: Path to PDB or MMTF file. Default is ``None``.
     :param pdb_code: PDB accesion code. Default is ``None``.
     :param uniprot_id: UniProt ID. Default is ``None``.
     :param chain_selection: Selection of chains to include (e.g. ``"ABC"``) or
@@ -141,11 +145,11 @@ def protein_to_pyg(
     """
 
     # Get ID
-    if pdb_path is not None:
+    if path is not None:
         id = (
-            pdb_path.split("/")[-1] + "_" + chain_selection
+            path.split("/")[-1] + "_" + chain_selection
             if chain_selection != "all"
-            else pdb_path
+            else path
         )
     elif pdb_code is not None:
         id = (
@@ -161,7 +165,7 @@ def protein_to_pyg(
         )
 
     df = read_pdb_to_dataframe(
-        pdb_path=pdb_path,
+        path=path,
         pdb_code=pdb_code,
         uniprot_id=uniprot_id,
         model_index=model_index,
diff --git a/graphein/protein/utils.py b/graphein/protein/utils.py
index 45d61e9e..79d381d9 100644
--- a/graphein/protein/utils.py
+++ b/graphein/protein/utils.py
@@ -224,16 +224,16 @@ def download_pdb(
     return out_dir / f"{pdb_code}{extension}"
 
 
-def get_protein_name_from_filename(pdb_path: str) -> str:
+def get_protein_name_from_filename(path: str) -> str:
     """
-    Extracts a filename from a ``pdb_path``
+    Extracts a filename from a ``path``
 
-    :param pdb_path: Path to extract filename from.
-    :type pdb_path: str
+    :param path: Path to extract filename from.
+    :type path: str
     :return: file name.
     :rtype: str
     """
-    _, tail = os.path.split(pdb_path)
+    _, tail = os.path.split(path)
     tail = os.path.splitext(tail)[0]
     return tail
 
diff --git a/graphein/rna/graphs.py b/graphein/rna/graphs.py
index 17d2a968..c81604ee 100644
--- a/graphein/rna/graphs.py
+++ b/graphein/rna/graphs.py
@@ -7,7 +7,8 @@
 # Code Repository: https://github.com/a-r-j/graphein
 # This submodule is heavily inspired by: https://github.com/emalgorithm/rna-design/blob/aec77a18abe4850958d6736ec185a6f8cbfdf20c/src/util.py#L9
 
-from typing import Callable, List, Optional
+import os
+from typing import Callable, List, Optional, Union
 
 import networkx as nx
 from loguru import logger as log
@@ -94,7 +95,7 @@ def validate_dotbracket(db: str):
 def construct_rna_graph_3d(
     config: Optional[RNAGraphConfig] = None,
     name: Optional[str] = None,
-    pdb_path: Optional[str] = None,
+    path: Optional[Union[str, os.PathLike]] = None,
     pdb_code: Optional[str] = None,
     chain_selection: str = "all",
     model_index: int = 1,
@@ -105,7 +106,7 @@ def construct_rna_graph_3d(
     graph_annotation_funcs: Optional[List[Callable]] = None,
 ) -> nx.Graph:
     """
-    Constructs RNA structure graph from a ``pdb_code`` or ``pdb_path``.
+    Constructs RNA structure graph from a ``pdb_code`` or ``path``.
 
     Users can provide a :class:`~graphein.rna.config.RNAGraphConfig`
     object to specify construction parameters.
@@ -116,9 +117,9 @@ def construct_rna_graph_3d(
     :param config: :class:`~graphein.rna.config.RNAGraphConfig` object. If
         ``None``, defaults to config in ``graphein.rna.config``.
     :type config: graphein.protein.config.RNAGraphConfig, optional
-    :param pdb_path: Path to ``pdb_file`` to build graph from. Default is
+    :param path: Path to PDB or MMTF to build graph from. Default is
         ``None``.
-    :type pdb_path: str, optional
+    :type path: str, optional
     :param pdb_code: 4-character PDB accession pdb_code to build graph from.
         Default is ``None``.
     :type pdb_code: str, optional
@@ -152,8 +153,8 @@ def construct_rna_graph_3d(
         config = RNAGraphConfig()
 
     # Get name from pdb_file is no pdb_code is provided
-    if pdb_path and (pdb_code is None):
-        pdb_code = gp.get_protein_name_from_filename(pdb_path)
+    if path and (pdb_code is None):
+        pdb_code = gp.get_protein_name_from_filename(path)
 
     # If config params are provided, overwrite them
     config.rna_df_processing_functions = (
@@ -183,7 +184,7 @@ def construct_rna_graph_3d(
     )
 
     raw_df = gp.read_pdb_to_dataframe(
-        pdb_path,
+        path,
         pdb_code,
         model_index=model_index,
     )
@@ -203,7 +204,7 @@ def construct_rna_graph_3d(
         raw_pdb_df=raw_df,
         name=name,
         pdb_code=pdb_code,
-        pdb_path=pdb_path,
+        path=path,
         granularity=config.granularity,
     )
     # Add nodes to graph
@@ -343,7 +344,7 @@ def construct_graph(
     use_nussinov: bool = False,
     name: Optional[str] = None,
     config: Optional[RNAGraphConfig] = None,
-    pdb_path: Optional[str] = None,
+    path: Optional[str] = None,
     pdb_code: Optional[str] = None,
     chain_selection: str = "all",
     rna_df_processing_funcs: Optional[List[Callable]] = None,
@@ -354,11 +355,11 @@ def construct_graph(
 ) -> nx.Graph:
     # TODO Docstring
 
-    if pdb_path is not None or pdb_code is not None:
+    if path is not None or pdb_code is not None:
         return construct_rna_graph_3d(
             config=config,
             name=name,
-            pdb_path=pdb_path,
+            path=path,
             pdb_code=pdb_code,
             chain_selection=chain_selection,
             rna_df_processing_funcs=rna_df_processing_funcs,
diff --git a/notebooks/alphafold_protein_graph_tutorial.ipynb b/notebooks/alphafold_protein_graph_tutorial.ipynb
index 5f66a658..8119697e 100644
--- a/notebooks/alphafold_protein_graph_tutorial.ipynb
+++ b/notebooks/alphafold_protein_graph_tutorial.ipynb
@@ -95281,7 +95281,7 @@
     "c = gp.ProteinGraphConfig(granularity='CA')\n",
     "\n",
     "# Construct the graph!\n",
-    "g = gp.construct_graph(pdb_path=protein_path)\n",
+    "g = gp.construct_graph(path=protein_path)\n",
     "\n",
     "# Plot\n",
     "gp.plotly_protein_structure_graph(g, node_size_multiplier=0.5, colour_nodes_by=\"residue_name\")"
@@ -151803,7 +151803,7 @@
     "                                ]\n",
     ")\n",
     "\n",
-    "g = gp.construct_graph(pdb_path=protein_path, config=config)\n",
+    "g = gp.construct_graph(path=protein_path, config=config)\n",
     "\n",
     "gp.plotly_protein_structure_graph(g, colour_edges_by=\"kind\", colour_nodes_by=\"residue_name\", label_node_ids=False, node_size_multiplier=2, node_size_min=5)"
    ]
diff --git a/notebooks/rna_graph_tutorial.ipynb b/notebooks/rna_graph_tutorial.ipynb
index 05ce837f..576eba6f 100644
--- a/notebooks/rna_graph_tutorial.ipynb
+++ b/notebooks/rna_graph_tutorial.ipynb
@@ -610624,7 +610624,7 @@
    ],
    "source": [
     "# Reading it back in\n",
-    "g = gr.construct_graph(pdb_path=\"selection.pdb\")\n",
+    "g = gr.construct_graph(path=\"selection.pdb\")\n",
     "gr.plotly_rna_structure_graph(g, colour_nodes_by=\"seq_position\")"
    ]
   }


Test patch:
diff --git a/tests/protein/edges/test_distance.py b/tests/protein/edges/test_distance.py
index b1a855ef..1f92ac02 100644
--- a/tests/protein/edges/test_distance.py
+++ b/tests/protein/edges/test_distance.py
@@ -51,7 +51,7 @@ def generate_graph():
     """Generate PDB network.
     This is a helper function.
     """
-    return construct_graph(pdb_path=str(DATA_PATH))
+    return construct_graph(path=str(DATA_PATH))
 
 
 @pytest.fixture(scope="module")
@@ -187,7 +187,7 @@ def test_add_cation_pi_interactions(net):
 def test_add_peptide_bonds():
     file_path = Path(__file__).parent.parent / "test_data/4hhb.pdb"
     # Peptide bonds are default
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     for u, v in G.edges():
         assert abs(int(u.split(":")[2]) - int(v.split(":")[2])) == 1
@@ -203,7 +203,7 @@ def test_add_sequence_distance_edges():
                 partial(add_sequence_distance_edges, d=d)
             ]
         )
-        G = construct_graph(pdb_path=str(file_path), config=config)
+        G = construct_graph(path=str(file_path), config=config)
         for u, v in G.edges():
             assert abs(int(u.split(":")[2]) - int(v.split(":")[2])) == d
 
diff --git a/tests/protein/test_graphs.py b/tests/protein/test_graphs.py
index 177da920..a8c9a326 100644
--- a/tests/protein/test_graphs.py
+++ b/tests/protein/test_graphs.py
@@ -58,7 +58,7 @@ def generate_graph():
     """Generate PDB network.
     This is a helper function.
     """
-    return construct_graph(pdb_path=str(DATA_PATH))
+    return construct_graph(path=str(DATA_PATH))
 
 
 @pytest.fixture(scope="module")
@@ -89,7 +89,7 @@ def test_construct_graph():
     Uses 4hhb PDB file as an example test case.
     """
     file_path = Path(__file__).parent / "test_data" / "4hhb.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
     assert isinstance(G, nx.Graph)
     assert len(G) == 574
 
@@ -133,17 +133,17 @@ def test_construct_graph_with_dssp():
     )  # should download 6rew.pdb to pdb_dir
 
     assert g_pdb.graph["pdb_code"] == "6rew"
-    assert g_pdb.graph["pdb_path"] is None
+    assert g_pdb.graph["path"] is None
     assert g_pdb.graph["name"] == g_pdb.graph["pdb_code"]
     assert len(g_pdb.graph["dssp_df"]) == 1365
 
     file_path = str(
         Path(__file__).parent / "test_data" / "alphafold_structure.pdb"
     )
-    g_local = construct_graph(config=dssp_prot_config, pdb_path=file_path)
+    g_local = construct_graph(config=dssp_prot_config, path=file_path)
 
     assert g_local.graph["pdb_code"] is None
-    assert g_local.graph["pdb_path"] == file_path
+    assert g_local.graph["path"] == file_path
     assert g_local.graph["name"] == "alphafold_structure"
     assert len(g_local.graph["dssp_df"]) == 382
 
@@ -185,7 +185,7 @@ def test_chain_selection():
     Uses 4hhb PDB file as an example test case.
     """
     file_path = Path(__file__).parent / "test_data" / "4hhb.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     # Check default construction contains all chains
     assert G.graph["chain_ids"] == ["A", "B", "C", "D"]
@@ -194,7 +194,7 @@ def test_chain_selection():
         assert d["chain_id"] in ["A", "B", "C", "D"]
 
     # Check graph contains only chain selection
-    G = construct_graph(pdb_path=str(file_path), chain_selection="AD")
+    G = construct_graph(path=str(file_path), chain_selection="AD")
     assert G.graph["chain_ids"] == ["A", "D"]
     # Check nodes only contain residues from chain selection
     for n, d in G.nodes(data=True):
@@ -225,7 +225,7 @@ def test_intramolecular_edges():
         ]
     }
     config = ProteinGraphConfig(**edge_functions)
-    G = construct_graph(pdb_path=str(file_path), config=config)
+    G = construct_graph(path=str(file_path), config=config)
     # Todo complete
 """
 
@@ -258,7 +258,7 @@ def test_distance_edges():
         ]
     }
     config = ProteinGraphConfig(**edge_functions)
-    G = construct_graph(pdb_path=str(file_path), config=config)
+    G = construct_graph(path=str(file_path), config=config)
     assert G is not None
 
 
@@ -288,7 +288,7 @@ def test_node_features():
         "dssp_config": DSSPConfig(),
     }
     config = ProteinGraphConfig(**config_params)
-    G = construct_graph(pdb_path=str(file_path), config=config)
+    G = construct_graph(path=str(file_path), config=config)
 
     # Check for existence of features
     for _, d in G.nodes(data=True):
@@ -318,7 +318,7 @@ def test_sequence_features():
         ]
     }
     config = ProteinGraphConfig(**sequence_feature_functions)
-    G = construct_graph(pdb_path=str(file_path), config=config)
+    G = construct_graph(path=str(file_path), config=config)
 
     # Check for existence on sequence-based features as node-level features
     # for n, d in G.nodes(data=True):
@@ -494,7 +494,7 @@ def test_secondary_structure_graphs():
         graph_metadata_functions=[secondary_structure],
         dssp_config=DSSPConfig(),
     )
-    g = construct_graph(pdb_path=str(file_path), config=config)
+    g = construct_graph(path=str(file_path), config=config)
 
     h = compute_secondary_structure_graph(g, remove_non_ss=False)
     # Check number of residues preserved
@@ -532,7 +532,7 @@ def test_chain_graph():
         graph_metadata_functions=[secondary_structure],
         dssp_config=DSSPConfig(),
     )
-    g = construct_graph(pdb_path=str(file_path), config=config)
+    g = construct_graph(path=str(file_path), config=config)
     h = compute_chain_graph(g)
     assert len(h.edges) == len(g.edges), "Number of edges do not match"
 
diff --git a/tests/protein/test_subgraphs.py b/tests/protein/test_subgraphs.py
index d3ddc65b..a293ad75 100644
--- a/tests/protein/test_subgraphs.py
+++ b/tests/protein/test_subgraphs.py
@@ -43,7 +43,7 @@ def test_node_list_subgraphing():
     file_path = Path(__file__).parent / "test_data/4hhb.pdb"
     NODE_LIST = ["C:ALA:28", "C:ARG:31", "D:LEU:75", "A:THR:38"]
 
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     g = extract_subgraph_from_node_list(G, NODE_LIST, filter_dataframe=True)
 
@@ -88,7 +88,7 @@ def test_node_list_subgraphing():
 def test_extract_subgraph_from_atom_types():
     """Tests subgraph extraction from a list of allowed atom types"""
     file_path = Path(__file__).parent / "test_data/4hhb.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     ATOM_TYPES = ["CA"]
     g = extract_subgraph_from_atom_types(G, ATOM_TYPES, filter_dataframe=True)
@@ -104,7 +104,7 @@ def test_extract_subgraph_from_residue_types():
     SERINES = 32
     GLYCINES = 40
 
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     g = extract_subgraph_from_residue_types(
         G, RESIDUE_TYPES, filter_dataframe=True
@@ -165,7 +165,7 @@ def test_extract_subgraph_from_residue_types():
 def test_extract_subgraph_from_point():
     """Tests subgraph extraction from a spherical selection."""
     file_path = Path(__file__).parent / "test_data/4hhb.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     POINT = np.array([0.0, 0.0, 0.0])
     RADIUS = 10
@@ -197,7 +197,7 @@ def test_extract_subgraph_from_point():
 def test_extract_subgraph_from_chains():
     """Tests subgraph extraction from chains."""
     file_path = Path(__file__).parent / "test_data/4hhb.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     CHAINS = ["A", "C"]
     s_g = extract_subgraph_from_chains(G, CHAINS, filter_dataframe=True)
@@ -232,7 +232,7 @@ def test_extract_subgraph_from_chains():
 def test_extract_subgraph_from_sequence_position():
     """Tests subgraph extraction from sequence position."""
     file_path = Path(__file__).parent / "test_data/4hhb.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     SEQ_POS = list(range(1, 50, 2))
 
@@ -272,7 +272,7 @@ def test_extract_subgraph_from_bond_type():
     config = ProteinGraphConfig(
         edge_construction_functions=[add_peptide_bonds, add_ionic_interactions]
     )
-    G = construct_graph(pdb_path=str(file_path))  # , config=config)
+    G = construct_graph(path=str(file_path))  # , config=config)
 
     BOND_TYPES = ["ionic"]
 
@@ -304,7 +304,7 @@ def test_extract_subgraph_from_bond_type():
 def test_extract_k_hop_subgraph():
     """Tests k-hop subgraph extraction."""
     file_path = Path(__file__).parent / "test_data/4hhb.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     CENTRAL_NODE = "B:SER:49"
     K = 1
@@ -325,7 +325,7 @@ def test_surface_subgraph():
     config = ProteinGraphConfig(
         graph_metadata_functions=[rsa], dssp_config=DSSPConfig()
     )
-    G = construct_graph(pdb_path=str(file_path), config=config)
+    G = construct_graph(path=str(file_path), config=config)
 
     RSA_THRESHOLD: float = 0.2
     s_g = extract_surface_subgraph(G, RSA_THRESHOLD, filter_dataframe=True)
@@ -346,7 +346,7 @@ def test_secondary_structure_subgraph():
         graph_metadata_functions=[secondary_structure],
         dssp_config=DSSPConfig(),
     )
-    G = construct_graph(pdb_path=str(file_path), config=config)
+    G = construct_graph(path=str(file_path), config=config)
 
     SS_ELEMENTS: List[str] = ["H"]
     s_g = extract_subgraph_from_secondary_structure(
@@ -369,7 +369,7 @@ def test_successful_pickle():
         graph_metadata_functions=[secondary_structure],
         dssp_config=DSSPConfig(),
     )
-    G = construct_graph(pdb_path=str(file_path), config=config)
+    G = construct_graph(path=str(file_path), config=config)
     s_g = extract_subgraph_from_residue_types(
         G,
         residue_types=["ALA", "SER", "MET"],
diff --git a/tests/protein/test_utils.py b/tests/protein/test_utils.py
index fded2374..022985d9 100644
--- a/tests/protein/test_utils.py
+++ b/tests/protein/test_utils.py
@@ -32,7 +32,7 @@ def test_save_graph_to_pdb():
         a.drop(["line_idx"], axis=1),
         g.graph["pdb_df"].drop(["line_idx", "node_id", "residue_id"], axis=1),
     )
-    h = construct_graph(pdb_path="/tmp/test_graph.pdb")
+    h = construct_graph(path="/tmp/test_graph.pdb")
 
     # We check for isomorphism rather than equality as array features are not
     # comparable
@@ -55,7 +55,7 @@ def test_save_pdb_df_to_pdb():
 
     # Now check for raw, unprocessed DF
     save_pdb_df_to_pdb(g.graph["raw_pdb_df"], "/tmp/test_pdb.pdb")
-    h = construct_graph(pdb_path="/tmp/test_pdb.pdb")
+    h = construct_graph(path="/tmp/test_pdb.pdb")
 
     # We check for isomorphism rather than equality as array features are not
     # comparable
diff --git a/tests/rna/test_graphs.py b/tests/rna/test_graphs.py
index 7ef82542..74c2babe 100644
--- a/tests/rna/test_graphs.py
+++ b/tests/rna/test_graphs.py
@@ -80,7 +80,7 @@ def test_pdb_rna_graph():
 
     assert "name" in g.graph.keys()
     assert "pdb_code" in g.graph.keys()
-    assert "pdb_path" in g.graph.keys()
+    assert "path" in g.graph.keys()
     assert "chain_ids" in g.graph.keys()
     assert "pdb_df" in g.graph.keys()
     assert "raw_pdb_df" in g.graph.keys()
@@ -92,7 +92,7 @@ def test_construct_graph():
     Uses 4hhb PDB file as an example test case.
     """
     file_path = Path(__file__).parent / "test_data/2jyf.pdb"
-    g = construct_graph(pdb_path=str(file_path))
+    g = construct_graph(path=str(file_path))
 
     for n, d in g.nodes(data=True):
         assert isinstance(n, str)
@@ -109,7 +109,7 @@ def test_construct_graph():
 
     assert "name" in g.graph.keys()
     assert "pdb_code" in g.graph.keys()
-    assert "pdb_path" in g.graph.keys()
+    assert "path" in g.graph.keys()
     assert "chain_ids" in g.graph.keys()
     assert "pdb_df" in g.graph.keys()
     assert "raw_pdb_df" in g.graph.keys()
diff --git a/tests/rna/test_subgraphs.py b/tests/rna/test_subgraphs.py
index 13ae21e1..7b4f4318 100644
--- a/tests/rna/test_subgraphs.py
+++ b/tests/rna/test_subgraphs.py
@@ -33,7 +33,7 @@ def test_node_list_subgraphing():
         "A:U:5:135:OP1",
     ]
 
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     g = extract_subgraph_from_node_list(G, NODE_LIST, filter_dataframe=True)
 
@@ -78,7 +78,7 @@ def test_node_list_subgraphing():
 def test_extract_subgraph_from_atom_types():
     """Tests subgraph extraction from a list of allowed atom types"""
     file_path = Path(__file__).parent / "test_data/2jyf.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     ATOM_TYPES = ["P"]
     g = extract_subgraph_from_atom_types(G, ATOM_TYPES, filter_dataframe=True)
@@ -94,7 +94,7 @@ def test_extract_subgraph_from_residue_types():
     A_COUNT = 484  # TODO
     G_COUNT = 592  # TODO
 
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     g = extract_subgraph_from_residue_types(
         G, RESIDUE_TYPES, filter_dataframe=True
@@ -151,7 +151,7 @@ def test_extract_subgraph_from_residue_types():
 def test_extract_subgraph_from_point():
     """Tests subgraph extraction from a spherical selection."""
     file_path = Path(__file__).parent / "test_data/2jyf.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     POINT = np.array([0.0, 0.0, 0.0])
     RADIUS = 10
@@ -183,7 +183,7 @@ def test_extract_subgraph_from_point():
 def test_extract_subgraph_from_chains():
     """Tests subgraph extraction from chains."""
     file_path = Path(__file__).parent / "test_data/2jyf.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     CHAINS = ["A", "C"]
     s_g = extract_subgraph_from_chains(G, CHAINS, filter_dataframe=True)
@@ -218,7 +218,7 @@ def test_extract_subgraph_from_chains():
 def test_extract_subgraph_from_sequence_position():
     """Tests subgraph extraction from sequence position."""
     file_path = Path(__file__).parent / "test_data/2jyf.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     SEQ_POS = list(range(1, 50, 2))
 
@@ -255,7 +255,7 @@ def test_extract_subgraph_from_sequence_position():
 def test_extract_subgraph_from_bond_type():
     """Tests subgraph extraction from bond type"""
     file_path = Path(__file__).parent / "test_data/2jyf.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     BOND_TYPES = ["covalent"]
 
@@ -287,7 +287,7 @@ def test_extract_subgraph_from_bond_type():
 def test_extract_k_hop_subgraph():
     """Tests k-hop subgraph extraction."""
     file_path = Path(__file__).parent / "test_data/2jyf.pdb"
-    G = construct_graph(pdb_path=str(file_path))
+    G = construct_graph(path=str(file_path))
 
     CENTRAL_NODE = "A:A:6:181:N6"
     K = 1
diff --git a/tests/testing/utils.py b/tests/testing/utils.py
index 00547a9e..3d1cd034 100644
--- a/tests/testing/utils.py
+++ b/tests/testing/utils.py
@@ -27,8 +27,8 @@
 def test_graph_isomorphism():
     """Test that the graph isomorphism is working."""
     file_path = Path(__file__).parent.parent / "protein" / "test_data/4hhb.pdb"
-    G_1 = construct_graph(pdb_path=str(file_path))
-    G_2 = construct_graph(pdb_path=str(file_path))
+    G_1 = construct_graph(path=str(file_path))
+    G_2 = construct_graph(path=str(file_path))
     assert graphs_isomorphic(G_1, G_2), "Graphs are not isomorphic."
     assert graphs_equal(G_1, G_2), "Graphs are not equal."
     assert nodes_equal(G_1, G_2), "Graphs do not contain the same nodes."
@@ -37,7 +37,7 @@ def test_graph_isomorphism():
     assert edge_data_equal(G_1, G_2), "Edge features differ."
 
     file_path = Path(__file__).parent.parent / "protein" / "test_data/1lds.pdb"
-    G_2 = construct_graph(pdb_path=str(file_path))
+    G_2 = construct_graph(path=str(file_path))
     assert not graphs_isomorphic(G_1, G_2), "Graphs not isomorphic."
     assert not graphs_equal(G_1, G_2), "Graphs are equal."
     assert not nodes_equal(G_1, G_2), "Graphs contain the same nodes."