From 3ee9b7506e4c321d413577db4dff2f7eecb28ad2 Mon Sep 17 00:00:00 2001
From: Alex Youngs <alexyoungs@hatchbed.com>
Date: Wed, 30 Oct 2024 17:30:26 -0500
Subject: [PATCH] Added install pkgs extension

---
 .gitignore                           |  3 +-
 README.md                            |  1 +
 setup.py                             |  1 +
 src/rocker/install_pkgs_extension.py | 43 ++++++++++++++++++++++++++++
 4 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 src/rocker/install_pkgs_extension.py

diff --git a/.gitignore b/.gitignore
index df2600f6..32befa2b 100644
--- a/.gitignore
+++ b/.gitignore
@@ -5,4 +5,5 @@ deb_dist
 *.pyc
 .coverage
 build
-rocker_venv
\ No newline at end of file
+rocker_venv
+.vscode
diff --git a/README.md b/README.md
index faea0499..19c0b3b9 100644
--- a/README.md
+++ b/README.md
@@ -37,6 +37,7 @@ You can get full details on the extensions from the main `rocker --help` command
 - home -- Mount the user's home directory into the container
 - pulse -- Mount pulse audio into the container
 - ssh -- Pass through ssh access to the container.
+- install_pkgs -- Install a list of additional packages not in the base image
 
 As well as access to many of the docker arguments as well such as `device`, `env`, `volume`, `name`, `network`, `ipc`, and `privileged`.
 
diff --git a/setup.py b/setup.py
index 4d669832..3305afd4 100644
--- a/setup.py
+++ b/setup.py
@@ -55,6 +55,7 @@
             'home = rocker.extensions:HomeDir',
             'hostname = rocker.extensions:Hostname',
             'ipc = rocker.extensions:Ipc',
+            'install_pkgs = rocker.install_pkgs_extension:InstallPkgs',
             'name = rocker.extensions:Name',
             'network = rocker.extensions:Network',
             'nvidia = rocker.nvidia_extension:Nvidia',
diff --git a/src/rocker/install_pkgs_extension.py b/src/rocker/install_pkgs_extension.py
new file mode 100644
index 00000000..da2556e2
--- /dev/null
+++ b/src/rocker/install_pkgs_extension.py
@@ -0,0 +1,43 @@
+import em
+import pkgutil
+from rocker.extensions import RockerExtension
+
+
+class InstallPkgs(RockerExtension):
+
+    name = 'install_pkgs'
+
+    @classmethod
+    def get_name(cls):
+        return cls.name
+
+    def __init__(self):
+        self._env_subs = None
+        self.name = InstallPkgs.get_name()
+
+    def precondition_environment(self, cli_args):
+        pass
+
+    def validate_environment(self, cli_args):
+        pass
+
+    def get_preamble(self, cli_args):
+        return ''
+
+    def get_snippet(self, cli_args):
+        pkgs = set(cli_args['install_pkgs'])
+        args = {'packages': list(pkgs)}
+
+        snippet = pkgutil.get_data(
+            'rocker', 'templates/{}_snippet.Dockerfile.em'.format(self.name)).decode('utf-8')
+
+        return em.expand(snippet, args)
+
+    @staticmethod
+    def register_arguments(parser, defaults={}):
+        parser.add_argument('--install-pkgs',
+                            nargs='+',
+                            help='Installs specified packages in container')
+
+        # todo add argument to install common development packages by category 
+        # (e.g dev, debug, viz, sim, etc) that are not in the base image