From ce697cdfa8776187760033cc3d46151c010d9991 Mon Sep 17 00:00:00 2001
From: Martin Weinelt <hexa@darmstadt.ccc.de>
Date: Thu, 7 Apr 2022 05:24:06 +0200
Subject: [PATCH] home-assistant: allow requirements parser to deal with extras

---
 .../home-assistant/parse-requirements.py      | 56 ++++++++++++++++---
 1 file changed, 47 insertions(+), 9 deletions(-)

diff --git a/pkgs/servers/home-assistant/parse-requirements.py b/pkgs/servers/home-assistant/parse-requirements.py
index 0e3d5d5553e7..c9fc53401ef4 100755
--- a/pkgs/servers/home-assistant/parse-requirements.py
+++ b/pkgs/servers/home-assistant/parse-requirements.py
@@ -98,13 +98,37 @@ def get_reqs(components: Dict[str, Dict[str, Any]], component: str, processed: S
     return requirements
 
 
+def repository_root() -> str:
+    return os.path.abspath(sys.argv[0] + "/../../../..")
+
+
+# For a package attribute and and an extra, check if the package exposes it via passthru.extras-require
+def has_extra(package: str, extra: str):
+    cmd = [
+        "nix-instantiate",
+        repository_root(),
+        "-A",
+        f"{package}.extras-require.{extra}",
+    ]
+    try:
+        subprocess.run(
+            cmd,
+            check=True,
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+    except subprocess.CalledProcessError:
+        return False
+    return True
+
+
 def dump_packages() -> Dict[str, Dict[str, str]]:
     # Store a JSON dump of Nixpkgs' python3Packages
     output = subprocess.check_output(
         [
             "nix-env",
             "-f",
-            os.path.dirname(sys.argv[0]) + "/../../..",
+            repository_root(),
             "-qa",
             "-A",
             PKG_SET,
@@ -158,6 +182,7 @@ def main() -> None:
     outdated = {}
     for component in sorted(components.keys()):
         attr_paths = []
+        extra_attrs = []
         missing_reqs = []
         reqs = sorted(get_reqs(components, component, set()))
         for req in reqs:
@@ -165,9 +190,10 @@ def main() -> None:
             # Therefore, if there's a "#" in the line, only take the part after it
             req = req[req.find("#") + 1 :]
             name, required_version = req.split("==", maxsplit=1)
-            # Remove extra_require from name, e.g. samsungctl instead of
-            # samsungctl[websocket]
+            # Split package name and extra requires
+            extras = []
             if name.endswith("]"):
+                extras = name[name.find("[")+1:name.find("]")].split(",")
                 name = name[:name.find("[")]
             attr_path = name_to_attr_path(name, packages)
             if our_version := get_pkg_version(name, packages):
@@ -178,11 +204,20 @@ def main() -> None:
                     }
             if attr_path is not None:
                 # Add attribute path without "python3Packages." prefix
-                attr_paths.append(attr_path[len(PKG_SET + ".") :])
+                pname = attr_path[len(PKG_SET + "."):]
+                attr_paths.append(pname)
+                for extra in extras:
+                    # Check if package advertises extra requirements
+                    extra_attr = f"{pname}.extras-require.{extra}"
+                    if has_extra(attr_path, extra):
+                        extra_attrs.append(extra_attr)
+                    else:
+                        missing_reqs.append(extra_attr)
+
             else:
                 missing_reqs.append(name)
         else:
-            build_inputs[component] = (attr_paths, missing_reqs)
+            build_inputs[component] = (attr_paths, extra_attrs, missing_reqs)
 
     with open(os.path.dirname(sys.argv[0]) + "/component-packages.nix", "w") as f:
         f.write("# Generated by parse-requirements.py\n")
@@ -191,11 +226,14 @@ def main() -> None:
         f.write(f'  version = "{version}";\n')
         f.write("  components = {\n")
         for component, deps in build_inputs.items():
-            available, missing = deps
+            available, extras, missing = deps
             f.write(f'    "{component}" = ps: with ps; [')
             if available:
-                f.write(" " + " ".join(available))
-            f.write(" ];")
+                f.write("\n      " + "\n      ".join(available))
+            f.write("\n    ]")
+            if extras:
+                f.write("\n    ++ " + "\n    ++ ".join(extras))
+            f.write(";")
             if len(missing) > 0:
                 f.write(f" # missing inputs: {' '.join(missing)}")
             f.write("\n")
@@ -203,7 +241,7 @@ def main() -> None:
         f.write("  # components listed in tests/components for which all dependencies are packaged\n")
         f.write("  supportedComponentsWithTests = [\n")
         for component, deps in build_inputs.items():
-            available, missing = deps
+            available, extras, missing = deps
             if len(missing) == 0 and component in components_with_tests:
                 f.write(f'    "{component}"' + "\n")
         f.write("  ];\n")