iteration 3

.github/workflows/scripts/merge_requirements.py

@@ -2,7 +2,7 @@
 
 import os
 from pathlib import Path
-from typing import List, TextIO
+from typing import Dict, Iterable, List, TextIO, Tuple
 
 from packaging.markers import Marker
 from packaging.requirements import Requirement
@@ -57,22 +57,28 @@ def get_requirements(fp: TextIO) -> List[RequirementData]:
     return requirements
 
 
+def iter_envs(envs: Iterable[str]) -> Iterable[Tuple[str, str]]:
+    for env_name in envs:
+        platform, python_version = env_name.split("-", maxsplit=1)
+        yield (platform, python_version)
+
+
 names = ["base"]
 names.extend(file.stem for file in REQUIREMENTS_FOLDER.glob("extra-*.in"))
-base_requirements = []
+base_requirements: List[RequirementData] = []
 
 for name in names:
     # {req_data: {sys_platform: RequirementData}
-    input_data = {}
+    input_data: Dict[RequirementData, Dict[str, RequirementData]] = {}
     all_envs = set()
-    # all_platforms = set()
-    # all_python_versions = set()
+    all_platforms = set()
+    all_python_versions = set()
     for file in REQUIREMENTS_FOLDER.glob(f"*-{name}.txt"):
         platform_name, python_version, _ = file.stem.split("-", maxsplit=2)
         env_name = f"{platform_name}-{python_version}"
         all_envs.add(env_name)
-        # all_platforms.add(platform_name)
-        # all_python_versions.add(python_version)
+        all_platforms.add(platform_name)
+        all_python_versions.add(python_version)
         with file.open(encoding="utf-8") as fp:
             requirements = get_requirements(fp)
 
@@ -82,8 +88,12 @@ def get_requirements(fp: TextIO) -> List[RequirementData]:
 
     output = base_requirements if name == "base" else []
     for req, envs in input_data.items():
-        python_versions_per_platform = {}
-        platforms_per_python_version = {}
+        # {platform: [python_versions...]}
+        python_versions_per_platform: Dict[str, List[str]] = {}
+        # {python_version: [platforms...]}
+        platforms_per_python_version: Dict[str, List[str]] = {}
+        platforms = python_versions_per_platform.keys()
+        python_versions = platforms_per_python_version.keys()
         for env_name, other_req in envs.items():
             platform_name, python_version = env_name.split("-", maxsplit=1)
             python_versions_per_platform.setdefault(platform_name, []).append(python_version)
@@ -110,43 +120,79 @@ def get_requirements(fp: TextIO) -> List[RequirementData]:
             output.append(req)
             continue
 
+        # At this point I'm wondering why I didn't just go for
+        # a more generic boolean algebra simplification (sympy.simplify_logic())...
         if len(set(map(frozenset, python_versions_per_platform.values()))) == 1:
-            # same Python version set for all platforms, limit by platform only
-            if len(envs) < len(all_envs - envs.keys()):
+            # All platforms have the same Python version set.
+            # Let's generate the base Python version marker.
+            python_version_marker = " or ".join(
+                f"python_version == '{env_name}'" for env_name in envs
+                # f"python_version == '{python_version}'"
+                # for python_version in python_versions
+            )
+
+            if len(platforms) == len(all_platforms):
+                # requirement present on all platforms
+                env_marker = python_version_marker
+            elif len(platforms) < len(all_platforms - platforms):
+                # requirement present on less platforms than not
                 platform_marker = " or ".join(
-                    f"sys_platform == '{platform}'" for platform in envs
+                    f"sys_platform == '{env_name}'" for env_name in envs
+                    # f"sys_platform == '{platform}'" for platform in platforms
                 )
+                env_marker = f"({python_version_marker}) and ({platform_marker})"
             else:
+                # requirement present on more platforms than not
                 platform_marker = " and ".join(
-                    f"sys_platform != '{platform}'" for platform in all_envs - envs.keys()
+                    f"sys_platform != '{env_name}'" for env_name in all_envs - envs.keys()
+                    # f"sys_platform != '{platform}'" for platform in all_platforms - platforms
                 )
+                env_marker = f"({python_version_marker}) and ({platform_marker})"
         elif len(set(map(frozenset, platforms_per_python_version.values()))) == 1:
-            # same platform set for all python version, limit by Python version only
-            if len(envs) < len(all_envs - envs.keys()):
-                platform_marker = " or ".join(
-                    f"python_version == '{platform}'" for platform in envs
+            # All Python versions have the same Python version set.
+            # Let's generate the base platform marker
+            platform_marker = " or ".join(
+                f"sys_platform == '{env_name}'" for env_name in envs
+                # f"sys_platform == '{platform}'" for platform in platforms
+            )
+
+            if len(python_versions) == len(all_python_versions):
+                # requirement present on all Python versions
+                env_marker = platform_marker
+            elif len(python_versions) < len(all_python_versions - python_versions):
+                # requirement present on less Python versions than not
+                python_version_marker = " or ".join(
+                    f"python_version == '{env_name}'" for env_name in envs
+                    # f"python_version == '{python_version}'" for python_version in python_versions
                 )
+                env_marker = f"({platform_marker}) and ({python_version_marker})"
             else:
-                platform_marker = " and ".join(
-                    f"python_version != '{platform}'" for platform in all_envs - envs.keys()
+                # requirement present on more Python versions than not
+                python_version_marker = " and ".join(
+                    f"python_version != '{env_name}'" for env_name in all_envs - envs.keys()
+                    # f"python_version != '{python_version}'"
+                    # for python_version in all_python_versions - python_versions
                 )
+                env_marker = f"({platform_marker}) and ({python_version_marker})"
         else:
+            # Fallback to generic case.
+
             if len(envs) < len(all_envs - envs.keys()):
                 platform_marker = " or ".join(
-                    f"(sys_platform == '{platform}' and python_version == '{platform}')"
-                    for platform in envs
+                    f"(sys_platform == '{platform}' and python_version == '{python_version}')"
+                    for platform, python_version in iter_envs(envs)
                 )
             else:
                 platform_marker = " and ".join(
-                    f"(sys_platform != '{platform}' and python_version != '{platform}')"
-                    for platform in all_envs - envs.keys()
+                    f"(sys_platform != '{platform}' and python_version != '{python_version}')"
+                    for platform, python_version in iter_envs(all_envs - envs.keys())
                 )
 
 
         new_marker = (
-            f"({req.marker}) and ({platform_marker})"
+            f"({req.marker}) and ({env_marker})"
             if req.marker is not None
-            else platform_marker
+            else env_marker
         )
         req.marker = Marker(new_marker)
         if base_req is not None and base_req.marker == req.marker: