From 861374460f04e302e016aa1ee7e9a8423758d9da Mon Sep 17 00:00:00 2001
From: Edan Bainglass <edan.bainglass@gmail.com>
Date: Thu, 22 Aug 2024 20:08:13 +0000
Subject: [PATCH] Simplify CSS loading utility

---
 aiidalab_widgets_base/__init__.py      |  4 +++-
 aiidalab_widgets_base/utils/loaders.py | 28 +++-----------------------
 notebooks/test_misc.ipynb              | 15 +++++---------
 tests/test_loaders.py                  |  9 ++++++---
 4 files changed, 17 insertions(+), 39 deletions(-)

diff --git a/aiidalab_widgets_base/__init__.py b/aiidalab_widgets_base/__init__.py
index a9f4dec55..9984aebb9 100644
--- a/aiidalab_widgets_base/__init__.py
+++ b/aiidalab_widgets_base/__init__.py
@@ -40,6 +40,8 @@ def is_running_in_jupyter():
 
 
 if is_running_in_jupyter():
+    from pathlib import Path
+
     from aiida.manage import get_profile
     from IPython.display import HTML, display
 
@@ -52,7 +54,7 @@ def is_running_in_jupyter():
     from .static import styles
     from .utils.loaders import load_css_stylesheet
 
-    load_css_stylesheet(package=styles)
+    load_css_stylesheet(css_dir=Path(styles.__file__))
 
 
 from .computational_resources import (
diff --git a/aiidalab_widgets_base/utils/loaders.py b/aiidalab_widgets_base/utils/loaders.py
index 7af25818f..30fd7d0c3 100644
--- a/aiidalab_widgets_base/utils/loaders.py
+++ b/aiidalab_widgets_base/utils/loaders.py
@@ -1,44 +1,22 @@
 from __future__ import annotations
 
-from importlib.resources import Package, files
 from pathlib import Path
 
 from IPython.display import Javascript, display
 
 
-def load_css_stylesheet(
-    package: Package | None = None,
-    css_path: str | Path = "",
-    filename: str = "",
-):
+def load_css_stylesheet(css_dir: Path, filename: str = ""):
     """Load a CSS stylesheet from a package and inject it into the DOM.
 
     Parameters
     ----------
-    `package` : `Package`, optional
-        The package where the CSS file is located.
-    `css_path` : `str` | `Path`, optional
+    `css_dir` : `Path`
         The path to the folder where the CSS file is located.
     `filename` : `str`, optional
         The name of the CSS file to load.
         If not provided, all CSS files in the package/folder will be loaded.
     """
-    if package:
-        root = files(package)
-        filenames = (
-            [root / filename]
-            if filename
-            else [
-                root / path.name
-                for path in root.iterdir()
-                if path.is_file() and path.name.endswith(".css")
-            ]
-        )
-    elif css_path:
-        path = Path(css_path)
-        filenames = [path / filename] if filename else [*path.glob("*.css")]
-    else:
-        raise ValueError("Either `package` or `path` must be provided.")
+    filenames = [css_dir / filename] if filename else [*css_dir.glob("*.css")]
 
     for fn in filenames:
         stylesheet = fn.read_text()
diff --git a/notebooks/test_misc.ipynb b/notebooks/test_misc.ipynb
index 85504a7c7..0aad5fd86 100644
--- a/notebooks/test_misc.ipynb
+++ b/notebooks/test_misc.ipynb
@@ -1,14 +1,5 @@
 {
  "cells": [
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import ipywidgets as ipw"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -26,9 +17,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from pathlib import Path\n",
+    "\n",
     "from aiidalab_widgets_base.utils.loaders import load_css_stylesheet\n",
     "\n",
-    "load_css_stylesheet(css_path=\"../tests_notebooks/static/styles\")"
+    "load_css_stylesheet(css_dir=Path(\"../tests_notebooks/static/styles\"))"
    ]
   },
   {
@@ -37,6 +30,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "import ipywidgets as ipw\n",
+    "\n",
     "label = ipw.Label(\"Testing\")\n",
     "label.add_class(\"red-text\")\n",
     "display(label)"
diff --git a/tests/test_loaders.py b/tests/test_loaders.py
index 21caa9186..f83440400 100644
--- a/tests/test_loaders.py
+++ b/tests/test_loaders.py
@@ -1,8 +1,11 @@
+from pathlib import Path
+
+from aiidalab_widgets_base.static import styles
 from aiidalab_widgets_base.utils.loaders import load_css_stylesheet
 
 
 def test_load_css_stylesheet():
     """Test `load_css_stylesheet` function."""
-    package = "aiidalab_widgets_base.static.styles"
-    load_css_stylesheet(package=package, filename="global.css")
-    load_css_stylesheet(package=package)
+    css_dir = Path(styles.__file__)
+    load_css_stylesheet(css_dir=css_dir, filename="global.css")
+    load_css_stylesheet(css_dir=css_dir)