From 2a7192167d4f06006ac72382e9cdd989bb20ab14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20S=C3=B8rensen?= <joasoe@gmail.com>
Date: Sat, 23 Oct 2021 13:19:56 +0200
Subject: [PATCH] Use regex instead of partition to section package definition
 (#58278)

---
 script/hassfest/requirements.py     | 13 +++++++++----
 tests/hassfest/test_requirements.py |  7 +++++--
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/script/hassfest/requirements.py b/script/hassfest/requirements.py
index 4d111265b1e8..26cb834e4e29 100644
--- a/script/hassfest/requirements.py
+++ b/script/hassfest/requirements.py
@@ -22,7 +22,9 @@ from .model import Config, Integration
 IGNORE_PACKAGES = {
     commented.lower().replace("_", "-") for commented in COMMENT_REQUIREMENTS
 }
-PACKAGE_REGEX = re.compile(r"^(?:--.+\s)?([-_\.\w\d]+).*==.+$")
+PACKAGE_REGEX = re.compile(
+    r"^(?:--.+\s)?([-_\.\w\d\[\]]+)(==|>=|<=|~=|!=|<|>|===)*(.*)$"
+)
 PIP_REGEX = re.compile(r"^(--.+\s)?([-_\.\w\d]+.*(?:==|>=|<=|~=|!=|<|>|===)?.*$)")
 SUPPORTED_PYTHON_TUPLES = [
     REQUIRED_PYTHON_VER[:2],
@@ -84,16 +86,19 @@ def validate_requirements_format(integration: Integration) -> bool:
             )
             continue
 
-        pkg, sep, version = req.partition("==")
+        pkg, sep, version = PACKAGE_REGEX.match(req).groups()
 
-        if not sep and integration.core:
+        if integration.core and sep != "==":
             integration.add_error(
                 "requirements",
                 f'Requirement {req} need to be pinned "<pkg name>==<version>".',
             )
             continue
 
-        if AwesomeVersion(version).strategy == AwesomeVersionStrategy.UNKNOWN:
+        if (
+            version
+            and AwesomeVersion(version).strategy == AwesomeVersionStrategy.UNKNOWN
+        ):
             integration.add_error(
                 "requirements",
                 f"Unable to parse package version ({version}) for {pkg}.",
diff --git a/tests/hassfest/test_requirements.py b/tests/hassfest/test_requirements.py
index c65716d5d929..079e77f909b9 100644
--- a/tests/hassfest/test_requirements.py
+++ b/tests/hassfest/test_requirements.py
@@ -45,7 +45,7 @@ def test_validate_requirements_format_wrongly_pinned(integration: Integration):
 
 def test_validate_requirements_format_ignore_pin_for_custom(integration: Integration):
     """Test requirement ignore pinning for custom."""
-    integration.manifest["requirements"] = ["test_package>=1"]
+    integration.manifest["requirements"] = ["test_package>=1", "test_package"]
     integration.path = Path("")
     assert validate_requirements_format(integration)
     assert len(integration.errors) == 0
@@ -63,6 +63,9 @@ def test_validate_requirements_format_invalid_version(integration: Integration):
 
 def test_validate_requirements_format_successful(integration: Integration):
     """Test requirement with successful result."""
-    integration.manifest["requirements"] = ["test_package==1.2.3"]
+    integration.manifest["requirements"] = [
+        "test_package==1.2.3",
+        "test_package[async]==1.2.3",
+    ]
     assert validate_requirements_format(integration)
     assert len(integration.errors) == 0