Repository: ad12/meddlr (language: python)

Issue:
Github handler
Add a handler to download, cache, and map github files to path on the local machine. This is useful for annotation files, config files, etc.

### TODOs
- [x] Change annotation file caching to follow this api
- [x] Add support for config files
- [x] verify this works with the command line process of handling configuration files


Discussion / hints:

Fix patch:
diff --git a/meddlr/utils/__init__.py b/meddlr/utils/__init__.py
index 7dda0597..8ba7364d 100644
--- a/meddlr/utils/__init__.py
+++ b/meddlr/utils/__init__.py
@@ -1,1 +1,2 @@
 from . import cluster  # noqa: F401
+from . import path  # noqa: F401
diff --git a/meddlr/utils/cluster.py b/meddlr/utils/cluster.py
index 90a30ce3..4d164a53 100644
--- a/meddlr/utils/cluster.py
+++ b/meddlr/utils/cluster.py
@@ -2,25 +2,16 @@
 
 DO NOT MOVE THIS FILE.
 """
-import logging
 import os
 import re
-import shutil
 import socket
-from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Sequence, Union
 
 import yaml
-from iopath.common.file_io import PathHandler
 
-import meddlr as mr
 from meddlr.utils import env
 
-# Path to the repository directory.
-# TODO: make this cleaner
-_REPO_DIR = os.path.join(os.path.dirname(__file__), "../..")
 _PATH_MANAGER = env.get_path_manager()
-_LOGGER = logging.getLogger(__name__)
 
 __all__ = ["Cluster", "set_cluster"]
 
@@ -296,130 +287,3 @@ def set_cluster(cluster: Union[str, Cluster] = None):
 
 _UNKNOWN = Cluster("UNKNOWN", [])  # Unknown cluster
 _CLUSTER = Cluster.cluster()  # Working cluster
-
-
-class GeneralPathHandler(PathHandler, ABC):
-    PREFIX = ""
-
-    def _get_supported_prefixes(self):
-        return [self.PREFIX]
-
-    def _get_local_path(self, path: str, **kwargs):
-        name = path[len(self.PREFIX) :]
-        return os.path.join(self._root_dir(), name)
-
-    def _open(self, path, mode="r", **kwargs):
-        return _PATH_MANAGER.open(self._get_local_path(path), mode, **kwargs)
-
-    def _mkdirs(self, path: str, **kwargs):
-        os.makedirs(self._get_local_path(path), exist_ok=True)
-
-    @abstractmethod
-    def _root_dir(self) -> str:
-        pass
-
-
-class DataHandler(GeneralPathHandler):
-    PREFIX = "data://"
-
-    def _root_dir(self):
-        return _CLUSTER.get_path("data_dir")
-
-
-class ResultsHandler(GeneralPathHandler):
-    PREFIX = "results://"
-
-    def _root_dir(self):
-        return _CLUSTER.get_path("results_dir")
-
-
-class CacheHandler(GeneralPathHandler):
-    PREFIX = "cache://"
-
-    def _root_dir(self):
-        return _CLUSTER.get_path("cache_dir")
-
-
-class AnnotationsHandler(GeneralPathHandler):
-    PREFIX = "ann://"
-
-    def _get_cache_path(self):
-        return os.path.join(
-            env.get_path_manager().get_local_path(_CLUSTER.cache_dir),
-            f"annotations/v{mr.__version__}",
-        )
-
-    def _root_dir(self):
-        local_path = os.path.abspath(os.path.join(_REPO_DIR, "annotations"))
-        if os.path.isdir(local_path):
-            return local_path
-
-        cache_path = self._get_cache_path()
-        if not os.path.isdir(cache_path):
-            self._download()
-        return cache_path
-
-    def _download(self):
-        """Downloads annotations from meddlr github.
-
-        Note:
-            This method downloads the full repository of the current
-            meddlr version and then copies the annotation files to
-            the respective directory. This may not be very efficient,
-            but it ensures that the annotation download process is
-            less onerous on the user.
-        """
-        _LOGGER.info("Downloading annotations...")
-        repo_path = download_repository()
-        cache_path = env.get_path_manager().get_local_path(self._get_cache_path())
-        os.makedirs(cache_path, exist_ok=True)
-        retval = os.system(f"cp -r {repo_path}/annotations {cache_path}/")
-        if retval != 0:
-            raise RuntimeError("Could not download annotations.")
-
-
-_PATH_MANAGER.register_handler(DataHandler())
-_PATH_MANAGER.register_handler(ResultsHandler())
-_PATH_MANAGER.register_handler(CacheHandler())
-_PATH_MANAGER.register_handler(AnnotationsHandler())
-
-
-def download_repository(version=None, path="cache://github-repo/{version}", force=False) -> str:
-    """Downloads the repository from GitHub.
-
-    Args:
-        version (str): The version to download. Defaults to the
-            version of the current codebase.
-        path (str): The path to download the repository to.
-            Defaults to the the cache directory.
-        force (bool): Whether to overwrite existing files.
-
-    Returns:
-        str: The path to the downloaded repository.
-    """
-    if version is None:
-        version = f"v{mr.__version__}"
-    path = str(path)
-    path = path.format(version=version)
-    path = env.get_path_manager().get_local_path(path)
-
-    dir_exists_and_not_empty = os.path.isdir(path) and os.listdir(path)
-    if dir_exists_and_not_empty:
-        if not force:
-            return path
-        shutil.rmtree(path)
-    os.makedirs(path, exist_ok=True)
-
-    retval = os.system(f"wget -O - {env.get_github_url()}/tarball/{version} | tar xz -C {path}")
-    if retval != 0:
-        raise RuntimeError("Could not download repository.")
-
-    folders = os.listdir(path)
-    assert len(folders) == 1
-    curr_path = os.path.join(path, folders[0])
-    retval = os.system(f"mv {curr_path}/* {path}")
-    if retval != 0:
-        raise RuntimeError("Could not download repository.")
-    shutil.rmtree(curr_path)
-
-    return path
diff --git a/meddlr/utils/path.py b/meddlr/utils/path.py
new file mode 100644
index 00000000..7c5af2f1
--- /dev/null
+++ b/meddlr/utils/path.py
@@ -0,0 +1,246 @@
+import logging
+import os
+import shutil
+from abc import ABC, abstractmethod
+from typing import Any, Optional
+
+from iopath.common.file_io import PathHandler
+
+import meddlr as mr
+from meddlr.utils import env
+from meddlr.utils.cluster import Cluster
+
+_REPO_DIR = os.path.join(os.path.dirname(__file__), "../..")
+_LOGGER = logging.getLogger(__name__)
+
+__all__ = ["GeneralPathHandler", "GithubHandler"]
+
+
+class GeneralPathHandler(PathHandler, ABC):
+    PREFIX = ""
+
+    def _get_supported_prefixes(self):
+        return [self.PREFIX]
+
+    def _get_local_path(self, path: str, **kwargs):
+        name = path[len(self.PREFIX) :]
+        return os.path.join(self._root_dir(), name)
+
+    def _open(self, path, mode="r", **kwargs):
+        return env.get_path_manager().open(self._get_local_path(path), mode, **kwargs)
+
+    def _mkdirs(self, path: str, **kwargs):
+        os.makedirs(self._get_local_path(path), exist_ok=True)
+
+    @abstractmethod
+    def _root_dir(self) -> str:
+        pass
+
+
+class DataHandler(GeneralPathHandler):
+    PREFIX = "data://"
+
+    def _root_dir(self):
+        return Cluster.working_cluster().get_path("data_dir")
+
+
+class ResultsHandler(GeneralPathHandler):
+    PREFIX = "results://"
+
+    def _root_dir(self):
+        return Cluster.working_cluster().get_path("results_dir")
+
+
+class CacheHandler(GeneralPathHandler):
+    PREFIX = "cache://"
+
+    def _root_dir(self):
+        return Cluster.working_cluster().get_path("cache_dir")
+
+
+class GithubHandler(GeneralPathHandler):
+    """Handler to download, cache, and match Github files to local folder structure.
+
+    Often, it may be useful to reference files from the meddlr Github repository.
+    This class (when used with :cls:`iopath.common.file_io.PathHandler`) fetches
+    and locally caches files from the Meddlr Github repository.
+
+    This class fetches files from the most recent tagged version in Meddlr.
+
+    Attributes:
+        github_url (str): The base URL for the Github repository.
+        default_branch_or_tag (str): The default branch or tag to download files from.
+
+    Examples:
+        >>> gh = GithubHandler(github_url="https://github.com/ad12/meddlr")
+        >>> gh.get_local_path("github://annotations")
+    """
+
+    PREFIX = "github://"
+
+    def __init__(
+        self,
+        github_url: str = None,
+        default_branch_or_tag: str = "main",
+        default_cache_path: str = None,
+        async_executor=None,
+    ) -> None:
+        if github_url is None:
+            github_url = env.get_github_url()
+        self.github_url = github_url
+        self.default_branch_or_tag = default_branch_or_tag
+        self.default_cache_path = default_cache_path
+        super().__init__(async_executor=async_executor)
+
+    def _root_dir(self):
+        return None
+
+    def _get_default_branch_or_tag(self):
+        return f"v{mr.__version__}"
+
+    def _get_default_cache_path(self):
+        return os.path.join(
+            env.get_path_manager().get_local_path(Cluster.working_cluster().cache_dir),
+            f"github-repo/v{mr.__version__}",
+        )
+
+    def download(self, cache_dir, branch_or_tag: str = None, force: bool = False):
+        if branch_or_tag is None:
+            branch_or_tag = self.default_branch_or_tag
+        if branch_or_tag is None:
+            branch_or_tag = self._get_default_branch_or_tag()
+        return download_github_repository(
+            url=self.github_url, branch_or_tag=branch_or_tag, cache_path=cache_dir, force=force
+        )
+
+    def _get_local_path(
+        self,
+        path: str,
+        branch_or_tag: str = None,
+        force: bool = False,
+        cache_dir: Optional[str] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Get local path to GitHub repository for the current meddlr version.
+
+        This implementation will only download the GitHub repository if it does
+        not exist. To force a download, set ``force=True``.
+
+        Args:
+            path (str): The relative file path in the GitHub repository.
+                Must start with ``'github://'``.
+            branch_or_tag (str): The Github branch or tag to download from.
+            force (bool, optional): If ``True``, force a download of the Github
+                repository. Defaults to ``False``.
+            cache_dir (str, optional): The local cache directory to use.
+
+        Returns:
+            str: The local path to the file/directory.
+        """
+        path = path[len(self.PREFIX) :]
+        self._check_kwargs(kwargs)
+        if cache_dir is None:
+            cache_dir = (
+                self.default_cache_path
+                if self.default_cache_path
+                else self._get_default_cache_path()
+            )
+
+        if force or not os.path.exists(cache_dir) or not os.listdir(cache_dir):
+            logger = logging.getLogger(__name__)
+            logger.info(
+                "Downloading github repository (tag/branch {}) from {}...".format(
+                    mr.__version__, self.github_url
+                )
+            )
+            self.download(cache_dir, force=force, branch_or_tag=branch_or_tag)
+            logger.info("Repository cached to {}".format(cache_dir))
+
+        expected_path = os.path.join(str(cache_dir), path)
+        if not os.path.exists(expected_path):
+            raise FileNotFoundError(
+                "Could not find {} in local Github cache: {}".format(path, cache_dir)
+            )
+        return expected_path
+
+
+class AnnotationsHandler(GeneralPathHandler):
+    PREFIX = "ann://"
+
+    def _get_cache_path(self):
+        return "github://annotations"
+
+    def _root_dir(self):
+        local_path = os.path.abspath(os.path.join(_REPO_DIR, "annotations"))
+        if os.path.isdir(local_path):
+            return local_path
+
+        cache_path = env.get_path_manager().get_local_path(self._get_cache_path())
+        if not os.path.isdir(cache_path):
+            self.download()
+        return cache_path
+
+    def download(self):
+        """Downloads annotations from meddlr github.
+
+        Note:
+            This method downloads the full repository of the current
+            meddlr version and then copies the annotation files to
+            the respective directory. This may not be very efficient,
+            but it ensures that the annotation download process is
+            less onerous on the user.
+        """
+        _LOGGER.info("Downloading annotations...")
+        repo_path = download_github_repository()
+        cache_path = env.get_path_manager().get_local_path(self._get_cache_path())
+        os.makedirs(cache_path, exist_ok=True)
+        retval = os.system(f"cp -r {repo_path}/annotations {cache_path}/")
+        if retval != 0:
+            raise RuntimeError("Could not download annotations.")
+
+
+def download_github_repository(url, cache_path, branch_or_tag="main", force=False) -> str:
+    """Downloads the repository from Github.
+
+    Args:
+        version (str): The version to download. Defaults to the
+            version of the current codebase.
+        path (str): The path to download the repository to.
+            Defaults to the the cache directory.
+        force (bool): Whether to overwrite existing files.
+
+    Returns:
+        str: The path to the downloaded repository.
+    """
+    url = f"{url}/tarball/{branch_or_tag}"
+    cache_path = str(cache_path)
+    cache_path = cache_path.format(branch_or_tag=branch_or_tag)
+
+    dir_exists_and_not_empty = os.path.isdir(cache_path) and os.listdir(cache_path)
+    if dir_exists_and_not_empty:
+        if not force:
+            return cache_path
+        shutil.rmtree(cache_path)
+    os.makedirs(cache_path, exist_ok=True)
+
+    retval = os.system(f"wget -O - {url} | tar xz -C {cache_path}")
+    if retval != 0:
+        raise RuntimeError("Could not download repository.")
+
+    folders = os.listdir(cache_path)
+    assert len(folders) == 1
+    curr_path = os.path.join(cache_path, folders[0])
+    retval = os.system(f"mv {curr_path}/* {cache_path}")
+    if retval != 0:
+        raise RuntimeError("Could not download repository.")
+    shutil.rmtree(curr_path)
+
+    return cache_path
+
+
+_path_manager = env.get_path_manager()
+_path_manager.register_handler(DataHandler())
+_path_manager.register_handler(ResultsHandler())
+_path_manager.register_handler(CacheHandler())
+_path_manager.register_handler(GithubHandler(env.get_github_url(), default_branch_or_tag=None))
+_path_manager.register_handler(AnnotationsHandler())
diff --git a/tools/train_net.py b/tools/train_net.py
index c89f7577..87321958 100644
--- a/tools/train_net.py
+++ b/tools/train_net.py
@@ -18,7 +18,7 @@
 from meddlr.config import get_cfg
 from meddlr.engine import DefaultTrainer, default_argument_parser, default_setup
 from meddlr.engine.defaults import init_wandb_run
-from meddlr.utils.env import supports_wandb
+from meddlr.utils import env
 
 
 def setup(args):
@@ -26,7 +26,7 @@ def setup(args):
     Create configs and perform basic setups.
     """
     cfg = get_cfg()
-    cfg.merge_from_file(args.config_file)
+    cfg.merge_from_file(env.get_path_manager().get_local_path(args.config_file))
     opts = args.opts
     if opts and opts[0] == "--":
         opts = opts[1:]
@@ -42,7 +42,7 @@ def setup(args):
     # Currently resuming with the same experiment id overwrites previous data.
     # So for now, even if you are resuming your experiment, it will be logged
     # as a separate run in W&B.
-    if supports_wandb():
+    if env.supports_wandb():
         init_wandb_run(cfg, resume=False, job_type="training", sync_tensorboard=True)
 
     return cfg


Test patch:
diff --git a/tests/utils/test_cluster.py b/tests/utils/test_cluster.py
index 7ba224d0..ddda8b5d 100644
--- a/tests/utils/test_cluster.py
+++ b/tests/utils/test_cluster.py
@@ -2,7 +2,7 @@
 import socket
 import unittest
 
-from meddlr.utils.cluster import Cluster, download_repository
+from meddlr.utils.cluster import Cluster
 
 from ..util import temp_env
 
@@ -73,9 +73,3 @@ def test_peristence(self):
         with self.assertRaises(KeyError):
             # Once deleted the cluster should not be found.
             cluster2 = cluster.from_config(name)
-
-
-def test_download_github(tmpdir):
-    download_dir = tmpdir.mkdir("download")
-    download_repository(version="main", path=download_dir)
-    assert os.path.isdir(download_dir.join("annotations").strpath)
diff --git a/tests/utils/test_path.py b/tests/utils/test_path.py
new file mode 100644
index 00000000..10d330fe
--- /dev/null
+++ b/tests/utils/test_path.py
@@ -0,0 +1,29 @@
+import os
+
+from meddlr.utils import env
+from meddlr.utils.path import GithubHandler, download_github_repository
+
+# def test_meddlr_path_manager():
+#     pm = env.get_path_manager()
+
+#     # Annotations
+#     assert os.path.isdir(pm.get_local_path("ann://"))
+
+#     # Github
+#     assert os.path.isdir(pm.get_local_path("github://"))
+
+
+def test_download_github(tmpdir):
+    download_dir = tmpdir.mkdir("download")
+    download_github_repository(env.get_github_url(), branch_or_tag="main", cache_path=download_dir)
+    assert os.path.isdir(download_dir.join("annotations").strpath)
+
+
+def test_github_handler(tmpdir):
+    download_dir = tmpdir.mkdir("download")
+    handler = GithubHandler(
+        env.get_github_url(), default_branch_or_tag="main", default_cache_path=download_dir / "main"
+    )
+    path = handler._get_local_path("github://annotations")
+    assert str(path) == os.path.join(download_dir.strpath, "main", "annotations")
+    assert os.path.isdir(path)