From 3e711f03f96a56bfb2ecfffa9a906cabea52aba9 Mon Sep 17 00:00:00 2001
From: Edgard Marx <edgard.marx@eccenca.com>
Date: Wed, 3 Apr 2024 10:43:16 +0200
Subject: [PATCH] first commit, adding library and cli

---
 src/imr-cli.py | 70 ++++++++++++++++++++++++++++++++++++++++++
 src/imr.py     | 83 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 153 insertions(+)
 create mode 100644 src/imr-cli.py
 create mode 100644 src/imr.py

diff --git a/src/imr-cli.py b/src/imr-cli.py
new file mode 100644
index 0000000..029f922
--- /dev/null
+++ b/src/imr-cli.py
@@ -0,0 +1,70 @@
+import click
+import yaml
+from os.path import expanduser
+from imr import IMRRemote, IMRLocal
+
+home = expanduser("~")
+imr_dir = home + "/.imr"
+imr_local : IMRLocal = IMRLocal(imr_dir)
+imr_remote : IMRRemote = None
+
+def loadParams():
+    with open(imr_dir + "/config.yaml", 'r') as stream:
+        imrConfig = yaml.safe_load(stream)
+
+@click.group()
+def cli():
+#    loadParams()
+    pass
+
+@cli.group()
+def local():
+    pass
+
+@local.command()
+def list():
+    for package in imr_local.list():
+        print(package)
+
+@local.command()
+@click.argument('package')
+@click.option('-v', '--version', type=str, default='latest', help='version of the model.', show_default=True)
+def remove(package, version):
+    list = local.rm(package, version)
+
+@cli.group()
+@click.argument('host')
+@click.argument('user')
+@click.argument('password')
+def remote(host, user, password):
+    global imr_remote
+    imr_remote = IMRRemote(host, user, password)
+
+@remote.command()
+def list():
+    list = imr_remote.list()
+    for p in list:
+        print(p)
+
+@remote.command()
+@click.argument('package')
+@click.option('-v', '--version', type=str, default='latest', help='version of the model.', show_default=True)
+def rm(package, version):
+    list = imr_remote.rm(package, version)
+
+@remote.command()
+@click.argument('model_directory')
+@click.argument('package')
+@click.option('-v', '--version', type=str, default='latest', help='version of the model.', show_default=True)
+def push(model_directory, package, version):
+    list = imr_remote.push(model_directory, package, version)
+
+@remote.command()
+@click.argument('package')
+@click.option('-d', '--dir', type=str, default=imr_dir, help='directory to pull the model in.', show_default=True)
+@click.option('-v', '--version', type=str, default='latest', help='version of the model.', show_default=True)
+def pull(package, dir, version):
+    list = imr_remote.pull(dir, package, version)
+
+if __name__ == '__main__':
+    cli()
\ No newline at end of file
diff --git a/src/imr.py b/src/imr.py
new file mode 100644
index 0000000..63b3c88
--- /dev/null
+++ b/src/imr.py
@@ -0,0 +1,83 @@
+# Inteligent Model Registry (imr)
+from artifactory import ArtifactoryPath
+from requests.auth import HTTPBasicAuth
+import shutil, os
+from os.path import expanduser
+from os import listdir
+
+class IMRRemote:
+
+    def __init__(self, repo:str, user:str, password:str):
+        self.repo = repo
+        self.user = user
+        self.password = password
+        self.home = expanduser("~")
+
+    def list(self):
+        packages = []
+        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
+        
+    def push(self,  directory: str, package: str, version = "latest"):
+        path = ArtifactoryPath(
+                    self.repo + "/" + package + "/" + version,
+                    auth=(self.user, self.password),
+                    auth_type=HTTPBasicAuth)
+        path.mkdir()
+        shutil.make_archive("model", 'zip', directory)
+        path.deploy_file("model.zip")
+
+    def pull(self, directory: str, package: str, version = "latest"):
+        path = ArtifactoryPath(
+                    self.repo + "/" + package + "/" + version + "/model.zip",
+                    auth=(self.user, self.password),
+                    auth_type=HTTPBasicAuth)
+        file_path = directory + "/" + package + "/" + version
+        if not os.path.exists(file_path):
+            os.makedirs(file_path)
+        with path.open() as fd, open(file_path + "/" + "model.zip", "wb") as out:
+            out.write(fd.read())
+        shutil.unpack_archive(file_path + "/" + "model.zip", file_path + "/model")
+    
+    def rm(self, package: str, version = "latest"):
+        artefact = self.repo + "/" + package
+        if version is not None:
+            artefact = self.repo + "/" + package + "/" + version
+        path = ArtifactoryPath(
+                    artefact,
+                    auth=(self.user, self.password),
+                    auth_type=HTTPBasicAuth)
+        if path.exists():
+            path.unlink()
+
+class IMRLocal:
+
+    def __init__(self, repo:str | None = None):
+        self.home = expanduser("~")
+        if(repo is None) :
+            self.repo = self.home + "/.imr"
+        else:
+            self.repo = repo
+        if not os.path.exists(self.repo):
+            os.makedirs(self.repo)
+
+    def list(self):
+        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
+        
+    def push(self,  directory: str, package: str, version = "latest"):
+        shutil.copytree(directory, self.repo + "/" + package + "/" + version)
+    
+    def rm(self, package: str, version = "latest"):
+        if version is None:
+            shutil.rmtree(self.repo + "/" + package)
+        else:
+            shutil.rmtree(self.repo + "/" + package + "/" + version)
\ No newline at end of file