diff --git a/imr/imr.py b/imr/imr.py index 494d0cc..d1a1f1f 100644 --- a/imr/imr.py +++ b/imr/imr.py @@ -8,6 +8,8 @@ class IMRRemote: + """Remote repository class.""" + def __init__(self, repo: str, user: str, password: str): self.repo = repo self.user = user @@ -15,13 +17,12 @@ def __init__(self, repo: str, user: str, password: str): self.home = str(Path.home()) def list(self) -> list: - packages = [] + """List models from the remote repository.""" path = ArtifactoryPath(self.repo, auth=(self.user, self.password), auth_type=HTTPBasicAuth) - for package in path.glob("*/*"): - packages.append(str(package).replace(self.repo + "/", "")) - return packages + return [str(package).replace(self.repo + "/", "") for package in path.glob("*/*")] - def push(self, directory: str, package: str, version :str="latest") -> None: + def push(self, directory: str, package: str, version: str = "latest") -> None: + """Push a local model to the remote repository.""" path = ArtifactoryPath( self.repo + "/" + package + "/" + version, auth=(self.user, self.password), @@ -31,7 +32,8 @@ def push(self, directory: str, package: str, version :str="latest") -> None: shutil.make_archive("model", "zip", directory) path.deploy_file("model.zip") - def pull(self, directory: str, package: str, version :str="latest") -> None: + def pull(self, directory: str, package: str, version: str = "latest") -> None: + """Pull a remote model to the local repository.""" path = ArtifactoryPath( self.repo + "/" + package + "/" + version + "/model.zip", auth=(self.user, self.password), @@ -44,7 +46,8 @@ def pull(self, directory: str, package: str, version :str="latest") -> None: out.write(fd.read()) shutil.unpack_archive(file_path + "/" + "model.zip", file_path + "/model") - def rm(self, package: str, version :str="latest") -> None: + def rm(self, package: str, version: str = "latest") -> None: + """Remove a model from the remote repository.""" artefact = self.repo + "/" + package if version is not None: artefact = self.repo + "/" + package + "/" + version @@ -54,6 +57,8 @@ def rm(self, package: str, version :str="latest") -> None: class IMRLocal: + """Local repository class.""" + def __init__(self, repo: str | None = None): self.home = Path.home() if repo is None: @@ -64,16 +69,19 @@ def __init__(self, repo: str | None = None): Path.mkdir(self.repo, parents=True) def list(self) -> list: - list = [] - for entry in os.walk(self.repo): - if len(entry[1]) == 0 and entry[0] not in self.repo: - list.append(entry[0].replace(self.repo + "/", "")) - return list + """List local models.""" + return [ + entry[0].replace(self.repo + "/", "") + for entry in os.walk(self.repo) + if len(entry[1]) == 0 and entry[0] not in self.repo + ] - def push(self, directory: str, package: str, version: str="latest") -> None: + def push(self, directory: str, package: str, version: str = "latest") -> None: + """Push a model in a directory to the local repository.""" shutil.copytree(directory, self.repo + "/" + package + "/" + version) - def rm(self, package: str, version :str="latest") -> None: + def rm(self, package: str, version: str = "latest") -> None: + """Remove a model from the local repository.""" if version is None: shutil.rmtree(self.repo + "/" + package) else: diff --git a/imr/imr-cli.py b/imr/imr_cli.py similarity index 72% rename from imr/imr-cli.py rename to imr/imr_cli.py index 2b02f71..93640f9 100644 --- a/imr/imr-cli.py +++ b/imr/imr_cli.py @@ -1,34 +1,40 @@ -from os.path import expanduser +from pathlib import Path +""" Path module. + +For handling local files and directory operations. +""" import click import yaml from imr import IMRLocal, IMRRemote -home = expanduser("~") +home = Path.home() imr_dir = home + "/.imr" imr_local: IMRLocal = IMRLocal(imr_dir) imr_remote: IMRRemote = None def loadParams() -> None: + """Load the default parameters from conf.yaml file.""" with open(imr_dir + "/config.yaml") as stream: imrConfig = yaml.safe_load(stream) @click.group() def cli() -> None: + """Get the cli command options.""" # loadParams() - pass @cli.group() def local() -> None: - pass + """Get local command options.""" -@local.command() +@local.command("list") def list_local() -> None: + """List local packages.""" for package in imr_local.list(): print(package) @@ -39,7 +45,8 @@ def list_local() -> None: "-v", "--version", type=str, default="latest", help="version of the model.", show_default=True ) def remove(package: str, version: str) -> None: - list = local.rm(package, version) + """Remove local packages.""" + local.rm(package, version) @cli.group() @@ -47,14 +54,16 @@ def remove(package: str, version: str) -> None: @click.argument("user") @click.argument("password") def remote(host: str, user: str, password: str) -> None: + """Get remote command cli options.""" global imr_remote imr_remote = IMRRemote(host, user, password) -@remote.command() +@remote.command("list") def list_remote() -> None: - list = imr_remote.list() - for p in list: + """List remote packages.""" + packages = imr_remote.list() + for p in packages: print(p) @@ -64,6 +73,7 @@ def list_remote() -> None: "-v", "--version", type=str, default="latest", help="version of the model.", show_default=True ) def rm(package: str, version: str) -> None: + """Remove remote package.""" imr_remote.rm(package, version) @@ -74,6 +84,7 @@ def rm(package: str, version: str) -> None: "-v", "--version", type=str, default="latest", help="version of the model.", show_default=True ) def push(model_directory: str, package: str, version: str) -> None: + """Push model to remote repository.""" imr_remote.push(model_directory, package, version) @@ -90,7 +101,8 @@ def push(model_directory: str, package: str, version: str) -> None: @click.option( "-v", "--version", type=str, default="latest", help="version of the model.", show_default=True ) -def pull(package :str, dir :str, version :str) -> None: +def pull(package: str, dir: str, version: str) -> None: + """Pull model from remote repository.""" imr_remote.pull(dir, package, version)