From 80218f33ebc0e717e119b6ca651293c80c01add5 Mon Sep 17 00:00:00 2001
From: jo <ljonas@riseup.net>
Date: Sun, 17 Jul 2022 18:51:58 +0200
Subject: [PATCH] chore: rewrite extract-requirements to python

Using python ast to guess requirements in the setup.py file.
---
 .pre-commit-config.yaml       |  3 +-
 analyzer/requirements.txt     |  2 +-
 api-client/requirements.txt   |  2 +-
 api/requirements.txt          | 14 ++++-----
 playout/requirements.txt      |  2 +-
 shared/requirements.txt       |  2 +-
 tools/extract-requirements.sh | 43 --------------------------
 tools/extract_requirements.py | 57 +++++++++++++++++++++++++++++++++++
 worker/requirements.txt       |  2 +-
 9 files changed, 70 insertions(+), 57 deletions(-)
 delete mode 100755 tools/extract-requirements.sh
 create mode 100755 tools/extract_requirements.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 0bba95a88..9e7c56ad0 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -23,7 +23,6 @@ repos:
         args: [--fix=lf]
       - id: trailing-whitespace
 
-      - id: requirements-txt-fixer
       - id: name-tests-test
         # TODO: Remove once the django api uses pytest
         exclude: ^(api.*)$
@@ -66,7 +65,7 @@ repos:
       - id: requirements.txt
         name: requirements.txt
         description: Generate requirements.txt
-        entry: tools/extract-requirements.sh dev
+        entry: tools/extract_requirements.py dev
         pass_filenames: false
         language: script
         files: setup.py$
diff --git a/analyzer/requirements.txt b/analyzer/requirements.txt
index 5f304d1fb..d2a71e781 100644
--- a/analyzer/requirements.txt
+++ b/analyzer/requirements.txt
@@ -1,5 +1,5 @@
-# This file is auto-generated by tools/extract-requirements.sh.
 # Please do not edit this file, edit the setup.py file!
+# This file is auto-generated by tools/extract_requirements.py.
 mutagen>=1.45.1
 pika>=1.0.0
 requests>=2.7.0
diff --git a/api-client/requirements.txt b/api-client/requirements.txt
index 5781df598..636f56a98 100644
--- a/api-client/requirements.txt
+++ b/api-client/requirements.txt
@@ -1,4 +1,4 @@
-# This file is auto-generated by tools/extract-requirements.sh.
 # Please do not edit this file, edit the setup.py file!
+# This file is auto-generated by tools/extract_requirements.py.
 python-dateutil>=2.7.0
 requests
diff --git a/api/requirements.txt b/api/requirements.txt
index f57101a13..a88135b20 100644
--- a/api/requirements.txt
+++ b/api/requirements.txt
@@ -1,8 +1,8 @@
-# This file is auto-generated by tools/extract-requirements.sh.
 # Please do not edit this file, edit the setup.py file!
-django<4.0,>=3.2.12
-django-filter<22.0,>=2.4.0
-djangorestframework<3.14,>=3.12.1
-drf-spectacular<0.23,>=0.22.1
-psycopg2<2.10,>=2.8.6
-requests<2.29,>=2.25.1
+# This file is auto-generated by tools/extract_requirements.py.
+django-filter>=2.4.0,<22.0
+django>=3.2.12,<4.0
+djangorestframework>=3.12.1,<3.14
+drf-spectacular>=0.22.1,<0.23
+psycopg2>=2.8.6,<2.10
+requests>=2.25.1,<2.29
diff --git a/playout/requirements.txt b/playout/requirements.txt
index f56b608b8..61a90e8bf 100644
--- a/playout/requirements.txt
+++ b/playout/requirements.txt
@@ -1,5 +1,5 @@
-# This file is auto-generated by tools/extract-requirements.sh.
 # Please do not edit this file, edit the setup.py file!
+# This file is auto-generated by tools/extract_requirements.py.
 amqplib
 defusedxml
 kombu
diff --git a/shared/requirements.txt b/shared/requirements.txt
index 8e1bce93f..f75199ca8 100644
--- a/shared/requirements.txt
+++ b/shared/requirements.txt
@@ -1,5 +1,5 @@
-# This file is auto-generated by tools/extract-requirements.sh.
 # Please do not edit this file, edit the setup.py file!
+# This file is auto-generated by tools/extract_requirements.py.
 click~=8.0.4
 loguru==0.6.0
 pydantic
diff --git a/tools/extract-requirements.sh b/tools/extract-requirements.sh
deleted file mode 100755
index 8bd97a780..000000000
--- a/tools/extract-requirements.sh
+++ /dev/null
@@ -1,43 +0,0 @@
-#!/usr/bin/env bash
-
-# Extract the dependencies from the setup.py files
-# and save the result to requirements.txt.
-#
-# You can filter any extra require by adding the name as argument.
-#
-# Examples:
-#   tools/extract-requirements.sh
-#   tools/extract-requirements.sh dev
-
-set -u
-
-error() {
-  echo >&2 "error: $*"
-  exit 1
-}
-
-command -v python3 > /dev/null || error "python3 command not found!"
-command -v sed > /dev/null || error "sed command not found!"
-
-for setup_path in */setup.py; do
-  path="$(dirname "$setup_path")"
-
-  # Build egg
-  python3 "$setup_path" egg_info > /dev/null 2>&1 || true
-  egg_path="$(echo "$path"/*.egg-info)"
-
-  # Remove entire extra section from require file
-  for arg in "$@"; do
-    sed --in-place "/^\[$arg\]/,/^\[/d" -- "$egg_path/requires.txt"
-  done
-
-  # Generate requirements.txt
-  cat << EOF > "$path/requirements.txt"
-# This file is auto-generated by tools/extract-requirements.sh.
-# Please do not edit this file, edit the setup.py file!
-EOF
-
-  cat -- "$egg_path/requires.txt" |
-    sed '/^$/d' |
-    LC_ALL=en_US.UTF-8 sort >> "$path/requirements.txt"
-done
diff --git a/tools/extract_requirements.py b/tools/extract_requirements.py
new file mode 100755
index 000000000..7d1eb887c
--- /dev/null
+++ b/tools/extract_requirements.py
@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+
+# Extract the dependencies from the setup.py files
+# and save the result to requirements.txt.
+#
+# You can filter any extra require by adding the name as argument.
+#
+# Examples:
+#   tools/extract_requirements.py
+#   tools/extract_requirements.py dev
+
+
+import ast
+from glob import glob
+from pathlib import Path
+from sys import argv
+
+
+class RemoveJoinedStr(ast.NodeTransformer):
+    def visit_JoinedStr(self, _node):  # pylint: disable=invalid-name
+        pass
+
+
+for setup in glob("*/setup.py"):
+    setup_path = Path(setup)
+    requirements_path = setup_path.parent / "requirements.txt"
+
+    lines = [
+        "# Please do not edit this file, edit the setup.py file!",
+        "# This file is auto-generated by tools/extract_requirements.py.",
+    ]
+
+    requires = []
+
+    for node in ast.walk(ast.parse(setup_path.read_text(encoding="utf-8"))):
+        if (
+            isinstance(node, ast.Expr)
+            and isinstance(node.value, ast.Call)
+            and isinstance(node.value.func, ast.Name)
+            and node.value.func.id == "setup"
+        ):
+
+            for keyword in node.value.keywords:
+                if keyword.arg == "install_requires":
+                    requires.extend(ast.literal_eval(keyword.value))
+
+                if keyword.arg == "extras_require":
+                    extras = ast.literal_eval(RemoveJoinedStr().visit(keyword.value))
+
+                    for key, values in extras.items():
+                        if key in argv:
+                            continue
+                        requires.extend(values)
+
+    lines.extend(sorted(requires))
+
+    requirements_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
diff --git a/worker/requirements.txt b/worker/requirements.txt
index 1a2d5b8e6..5f8b87db4 100644
--- a/worker/requirements.txt
+++ b/worker/requirements.txt
@@ -1,5 +1,5 @@
-# This file is auto-generated by tools/extract-requirements.sh.
 # Please do not edit this file, edit the setup.py file!
+# This file is auto-generated by tools/extract_requirements.py.
 celery==4.4.7
 kombu==4.6.11
 mutagen>=1.31.0