make match() on Enum into a top-level function in meta.py (pantsbuild…

…#8504)

### Problem

Importing via `from pants.util.collections import Enum` induces an indirection which isn't necessary for most users of `Enum` (the ones that don't use the `.match()` function).

### Solution

- Use the normal stdlib enum type.
- Make `match` into a top-level function in `meta.py`.
- Migrate the codebase.

### Result

An awkward import is resolved, and a new import is born! Complete credit for the solution goes to @jsirois in the `#engine` channel on our Slack.

build-support/bin/check_banned_imports.py

@@ -11,20 +11,7 @@
 
 def main() -> None:
   rust_files = find_files("src/rust/engine", extension=".rs")
-  python_files = find_files(
-    "src", "tests", "pants-plugins", "contrib", extension=".py"
-  )
 
-  check_banned_import(
-    python_files - {
-      "src/python/pants/base/hash_utils.py",
-      "src/python/pants/option/parser.py",
-      "src/python/pants/util/collections.py",
-      "tests/python/pants_test/base/test_hash_utils.py",
-    },
-    bad_import_regex=r"^from enum import.*Enum|^import enum$",
-    correct_import_message="from pants.util.collections import Enum",
-  )
   check_banned_import(
     rust_files,
     bad_import_regex=r"^use std::sync::.*(Mutex|RwLock)",

src/python/pants/backend/jvm/tasks/BUILD

@@ -573,6 +573,7 @@ python_library(
     'src/python/pants/java:util',
     'src/python/pants/process',
     'src/python/pants/task',
+    'src/python/pants/util:enums',
   ],
   tags = {"partially_type_checked"},
 )

src/python/pants/backend/jvm/tasks/jvm_compile/BUILD

@@ -62,6 +62,7 @@ python_library(
     'src/python/pants/util:collections',
     'src/python/pants/util:contextutil',
     'src/python/pants/util:dirutil',
+    'src/python/pants/util:enums',
     'src/python/pants/util:fileutil',
     'src/python/pants/util:memo',
   ],

src/python/pants/backend/jvm/tasks/jvm_compile/jvm_compile.py

@@ -3,6 +3,7 @@
 
 import functools
 import os
+from enum import Enum
 from multiprocessing import cpu_count
 from typing import Optional
 
@@ -39,7 +40,6 @@
 from pants.option.compiler_option_sets_mixin import CompilerOptionSetsMixin
 from pants.option.ranked_value import RankedValue
 from pants.reporting.reporting_utils import items_to_report_element
-from pants.util.collections import Enum
 from pants.util.contextutil import Timer, temporary_dir
 from pants.util.dirutil import (
   fast_relpath,
@@ -49,6 +49,7 @@
   safe_mkdir,
   safe_rmtree,
 )
+from pants.util.enums import match
 from pants.util.fileutil import create_size_estimators
 from pants.util.memo import memoized_method, memoized_property
 
@@ -958,7 +959,7 @@ def _get_jvm_distribution(self):
     # See: https://github.com/pantsbuild/pants/issues/6416 for covering using
     #      different jdks in remote builds.
     local_distribution = self._local_jvm_distribution()
-    return self.execution_strategy.match({
+    return match(self.execution_strategy, {
       self.ExecutionStrategy.subprocess: lambda: local_distribution,
       self.ExecutionStrategy.nailgun: lambda: local_distribution,
       self.ExecutionStrategy.hermetic: lambda: self._HermeticDistribution('.jdk', local_distribution),

src/python/pants/backend/jvm/tasks/jvm_compile/rsc/BUILD

@@ -21,6 +21,7 @@ python_library(
     'src/python/pants/util:collections',
     'src/python/pants/util:contextutil',
     'src/python/pants/util:dirutil',
+    'src/python/pants/util:enums',
     'src/python/pants/util:memo',
     'src/python/pants/util:strutil',
   ],

src/python/pants/backend/jvm/tasks/jvm_compile/rsc/rsc_compile.py

@@ -6,6 +6,7 @@
 import os
 import re
 from dataclasses import dataclass
+from enum import Enum
 
 from pants.backend.jvm.subsystems.dependency_context import DependencyContext  # noqa
 from pants.backend.jvm.subsystems.rsc import Rsc
@@ -30,9 +31,10 @@
 from pants.java.jar.jar_dependency import JarDependency
 from pants.reporting.reporting_utils import items_to_report_element
 from pants.task.scm_publish_mixin import Semver
-from pants.util.collections import Enum, assert_single_element
+from pants.util.collections import assert_single_element
 from pants.util.contextutil import Timer
 from pants.util.dirutil import fast_relpath, fast_relpath_optional, safe_mkdir
+from pants.util.enums import match
 from pants.util.memo import memoized_method, memoized_property
 from pants.util.strutil import safe_shlex_join
 
@@ -159,7 +161,7 @@ class JvmCompileWorkflowType(Enum):
   def _compiler_tags(self):
     return {
       f'{self.get_options().force_compiler_tag_prefix}:{workflow.value}': workflow
-      for workflow in self.JvmCompileWorkflowType.all_values()
+      for workflow in self.JvmCompileWorkflowType
     }
 
   @classmethod
@@ -171,12 +173,12 @@ def register_options(cls, register):
            'specified on the cli.')
 
     register('--workflow', type=cls.JvmCompileWorkflowType,
-      choices=cls.JvmCompileWorkflowType.all_values(),
+      choices=list(cls.JvmCompileWorkflowType),
       default=cls.JvmCompileWorkflowType.zinc_only, metavar='<workflow>',
       help='The default workflow to use to compile JVM targets. This is overriden on a per-target basis with the force-compiler-tag-prefix tag.', fingerprint=True)
 
     register('--scala-workflow-override', type=cls.JvmCompileWorkflowType,
-      choices=cls.JvmCompileWorkflowType.all_values(),
+      choices=list(cls.JvmCompileWorkflowType),
       default=None, metavar='<workflow_override>',
       help='Experimental option. The workflow to use to compile Scala targets, overriding the "workflow" option as well as any force-compiler-tag-prefix tags applied to targets. An example use case is to quickly turn off outlining workflows in case of errors.', fingerprint=True)
 
@@ -214,11 +216,11 @@ def register_options(cls, register):
         ),
       ],
       custom_rules=[
-        Shader.exclude_package('scala', recursive=True), 
-        Shader.exclude_package('xsbt', recursive=True), 
-        Shader.exclude_package('xsbti', recursive=True), 
-        # Unfortunately, is loaded reflectively by the compiler. 
-        Shader.exclude_package('org.apache.logging.log4j', recursive=True), 
+        Shader.exclude_package('scala', recursive=True),
+        Shader.exclude_package('xsbt', recursive=True),
+        Shader.exclude_package('xsbti', recursive=True),
+        # Unfortunately, is loaded reflectively by the compiler.
+        Shader.exclude_package('org.apache.logging.log4j', recursive=True),
       ]
     )
 
@@ -258,7 +260,7 @@ def _nailgunnable_combined_classpath(self):
 
   # Overrides the normal zinc compiler classpath, which only contains zinc.
   def get_zinc_compiler_classpath(self):
-    return self.execution_strategy.match({
+    return match(self.execution_strategy, {
       # NB: We must use the verbose version of super() here, possibly because of the lambda.
       self.ExecutionStrategy.hermetic: lambda: super(RscCompile, self).get_zinc_compiler_classpath(),
       self.ExecutionStrategy.subprocess: lambda: super(RscCompile, self).get_zinc_compiler_classpath(),
@@ -289,7 +291,7 @@ def confify(entries):
         rsc_cc.rsc_jar_file.hydrate_missing_directory_digest(classes_dir_snapshot.directory_digest)
 
       if rsc_cc.workflow is not None:
-        cp_entries = rsc_cc.workflow.match({
+        cp_entries = match(rsc_cc.workflow, {
           self.JvmCompileWorkflowType.zinc_only: lambda: confify([self._classpath_for_context(zinc_cc)]),
           self.JvmCompileWorkflowType.zinc_java: lambda: confify([self._classpath_for_context(zinc_cc)]),
           self.JvmCompileWorkflowType.rsc_and_zinc: lambda: confify([rsc_cc.rsc_jar_file]),
@@ -347,7 +349,7 @@ def _classify_target_compile_workflow(self, target):
 
   def _key_for_target_as_dep(self, target, workflow):
     # used for jobs that are either rsc jobs or zinc jobs run against rsc
-    return workflow.match({
+    return match(workflow, {
       self.JvmCompileWorkflowType.zinc_only: lambda: self._zinc_key_for_target(target, workflow),
       self.JvmCompileWorkflowType.zinc_java: lambda: self._zinc_key_for_target(target, workflow),
       self.JvmCompileWorkflowType.rsc_and_zinc: lambda: self._rsc_key_for_target(target),
@@ -361,7 +363,7 @@ def _outline_key_for_target(self, target):
     return f'outline({target.address.spec})'
 
   def _zinc_key_for_target(self, target, workflow):
-    return workflow.match({
+    return match(workflow, {
       self.JvmCompileWorkflowType.zinc_only: lambda: f'zinc[zinc-only]({target.address.spec})',
       self.JvmCompileWorkflowType.zinc_java: lambda: f'zinc[zinc-java]({target.address.spec})',
       self.JvmCompileWorkflowType.rsc_and_zinc: lambda: f'zinc[rsc-and-zinc]({target.address.spec})',
@@ -452,7 +454,7 @@ def nonhermetic_digest_classpath():
             classpath_abs_jdk = classpath_paths + self._jdk_libs_abs(distribution)
             return ((EMPTY_DIRECTORY_DIGEST), classpath_abs_jdk)
 
-          (input_digest, classpath_entry_paths) = self.execution_strategy.match({
+          (input_digest, classpath_entry_paths) = match(self.execution_strategy, {
             self.ExecutionStrategy.hermetic: hermetic_digest_classpath,
             self.ExecutionStrategy.subprocess: nonhermetic_digest_classpath,
             self.ExecutionStrategy.nailgun: nonhermetic_digest_classpath,
@@ -468,7 +470,7 @@ def nonhermetic_digest_classpath():
             ]
 
           target_sources = ctx.sources
-          
+
           # TODO: m.jar digests aren't found, so hermetic will fail.
           if use_youtline and not hermetic and self.get_options().zinc_outline:
             self._zinc_outline(ctx, classpath_paths, target_sources, youtline_args)
@@ -584,7 +586,7 @@ def record(k, v):
     record('execution_strategy', self.execution_strategy.value)
 
     # Create the cache doublecheck job.
-    workflow.match({
+    match(workflow, {
       self.JvmCompileWorkflowType.zinc_only: lambda: cache_doublecheck_jobs.append(
         make_cache_doublecheck_job(list(all_zinc_rsc_invalid_dep_keys(invalid_dependencies)))
       ),
@@ -601,7 +603,7 @@ def record(k, v):
 
     # Create the rsc job.
     # Currently, rsc only supports outlining scala.
-    workflow.match({
+    match(workflow, {
       self.JvmCompileWorkflowType.zinc_only: lambda: None,
       self.JvmCompileWorkflowType.zinc_java: lambda: None,
       self.JvmCompileWorkflowType.rsc_and_zinc: lambda: rsc_jobs.append(make_outline_job(compile_target, invalid_dependencies)),
@@ -612,7 +614,7 @@ def record(k, v):
     # - Scala zinc compile jobs depend on the results of running rsc on the scala target.
     # - Java zinc compile jobs depend on the zinc compiles of their dependencies, because we can't
     #   generate jars that make javac happy at this point.
-    workflow.match({
+    match(workflow, {
       # NB: zinc-only zinc jobs run zinc and depend on rsc and/or zinc compile outputs.
       self.JvmCompileWorkflowType.zinc_only: lambda: zinc_jobs.append(
         make_zinc_job(
@@ -725,7 +727,7 @@ def create_compile_context(self, target, target_workdir):
 
   def _runtool_hermetic(self, main, tool_name, distribution, input_digest, ctx):
     use_youtline = tool_name == 'scalac-outliner'
-    
+
     tool_classpath_abs = self._scalac_classpath if use_youtline else self._rsc_classpath
     tool_classpath = fast_relpath_collection(tool_classpath_abs)
 
@@ -897,7 +899,7 @@ def _runtool(self, distribution, input_digest, ctx, use_youtline):
       nailgun_classpath = self._nailgunnable_combined_classpath
 
     with self.context.new_workunit(tool_name) as wu:
-      return self.execution_strategy.match({
+      return match(self.execution_strategy, {
         self.ExecutionStrategy.hermetic: lambda: self._runtool_hermetic(
           main, tool_name, distribution, input_digest, ctx),
         self.ExecutionStrategy.subprocess: lambda: self._runtool_nonhermetic(

src/python/pants/backend/jvm/tasks/jvm_compile/zinc/BUILD

@@ -20,6 +20,7 @@ python_library(
     'src/python/pants/option',
     'src/python/pants/util:contextutil',
     'src/python/pants/util:dirutil',
+    'src/python/pants/util:enums',
     'src/python/pants/util:memo',
     'src/python/pants/util:strutil',
   ],

src/python/pants/backend/jvm/tasks/jvm_compile/zinc/zinc_compile.py

@@ -33,6 +33,7 @@
 from pants.engine.isolated_process import ExecuteProcessRequest
 from pants.util.contextutil import open_zip
 from pants.util.dirutil import fast_relpath
+from pants.util.enums import match
 from pants.util.memo import memoized_method, memoized_property
 from pants.util.meta import classproperty
 from pants.util.strutil import safe_shlex_join
@@ -364,7 +365,7 @@ def relative_to_exec_root(path):
     self.log_zinc_file(ctx.analysis_file)
     self.write_argsfile(ctx, zinc_args)
 
-    return self.execution_strategy.match({
+    return match(self.execution_strategy, {
       self.ExecutionStrategy.hermetic: lambda: self._compile_hermetic(
         jvm_options, ctx, classes_dir, jar_file, compiler_bridge_classpath_entry,
         dependency_classpath, scalac_classpath_entries),

src/python/pants/backend/jvm/tasks/nailgun_task.py

@@ -2,6 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 import os
+from enum import Enum
 
 from pants.backend.jvm.tasks.jvm_tool_task_mixin import JvmToolTaskMixin
 from pants.base.build_environment import get_buildroot
@@ -12,7 +13,6 @@
 from pants.java.nailgun_executor import NailgunExecutor, NailgunProcessGroup
 from pants.process.subprocess import Subprocess
 from pants.task.task import Task, TaskBase
-from pants.util.collections import Enum
 
 
 class NailgunTaskBase(JvmToolTaskMixin, TaskBase):

src/python/pants/backend/native/subsystems/BUILD

@@ -13,6 +13,7 @@ python_library(
     'src/python/pants/engine:selectors',
     'src/python/pants/subsystem',
     'src/python/pants/util:collections',
+    'src/python/pants/util:enums',
     'src/python/pants/util:memo',
     'src/python/pants/util:meta',
   ],

src/python/pants/backend/native/subsystems/binaries/BUILD

@@ -10,6 +10,7 @@ python_library(
     'src/python/pants/engine:rules',
     'src/python/pants/engine:platform',
     'src/python/pants/util:dirutil',
+    'src/python/pants/util:enums',
     'src/python/pants/util:memo',
   ],
   tags = {"partially_type_checked"},

src/python/pants/backend/native/subsystems/binaries/gcc.py

@@ -8,6 +8,7 @@
 from pants.binaries.binary_tool import NativeTool
 from pants.engine.platform import Platform
 from pants.engine.rules import rule
+from pants.util.enums import match
 from pants.util.memo import memoized_method, memoized_property
 
 
@@ -41,7 +42,7 @@ def path_entries(self):
 
   @memoized_method
   def _common_lib_dirs(self, platform):
-    lib64_tuples = platform.match({
+    lib64_tuples = match(platform, {
       Platform.darwin: [],
       Platform.linux: [('lib64',)]
     })

src/python/pants/backend/native/subsystems/binaries/llvm.py

@@ -10,6 +10,7 @@
 from pants.engine.platform import Platform
 from pants.engine.rules import RootRule, rule
 from pants.util.dirutil import is_readable_dir
+from pants.util.enums import match
 from pants.util.memo import memoized_method, memoized_property
 
 
@@ -81,7 +82,7 @@ def path_entries(self):
   def linker(self, platform: Platform) -> Linker:
     return Linker(
       path_entries=self.path_entries,
-      exe_filename=platform.match({
+      exe_filename=match(platform, {
         Platform.darwin: "ld64.lld",
         Platform.linux: "lld",
       }),

src/python/pants/backend/native/subsystems/native_build_step.py

@@ -2,12 +2,13 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
 from abc import abstractmethod
+from enum import Enum
 
 from pants.build_graph.mirrored_target_option_mixin import MirroredTargetOptionMixin
 from pants.engine.platform import Platform
 from pants.option.compiler_option_sets_mixin import CompilerOptionSetsMixin
 from pants.subsystem.subsystem import Subsystem
-from pants.util.collections import Enum
+from pants.util.enums import match
 from pants.util.memo import memoized_property
 from pants.util.meta import classproperty
 
@@ -37,7 +38,7 @@ def register_options(cls, register):
                   'for targets of this language.')
 
     register('--toolchain-variant', advanced=True,
-             default=Platform.current.match({
+             default=match(Platform.current, {
                Platform.darwin: ToolchainVariant.llvm,
                Platform.linux: ToolchainVariant.gnu,
              }),

src/python/pants/backend/native/subsystems/native_toolchain.py

@@ -22,6 +22,7 @@
 from pants.engine.rules import RootRule, rule
 from pants.engine.selectors import Get
 from pants.subsystem.subsystem import Subsystem
+from pants.util.enums import match
 from pants.util.memo import memoized_property
 
 
@@ -119,7 +120,7 @@ class LLVMCppToolchain:
 @rule
 async def select_libc_objects(platform: Platform, native_toolchain: NativeToolchain) -> LibcObjects:
   # We use lambdas here to avoid searching for libc on osx, where it will fail.
-  paths = platform.match({
+  paths = match(platform, {
     Platform.darwin: lambda: [],
     Platform.linux: lambda: native_toolchain._libc_dev.get_libc_objects(),
   })()

src/python/pants/backend/native/targets/BUILD

@@ -11,6 +11,7 @@ python_library(
     'src/python/pants/base:payload_field',
     'src/python/pants/build_graph',
     'src/python/pants/engine:platform',
+    'src/python/pants/util:enums',
     'src/python/pants/util:memo',
   ],
   tags = {"partially_type_checked"},

src/python/pants/backend/native/targets/native_artifact.py

@@ -7,6 +7,7 @@
 
 from pants.base.payload_field import PayloadField
 from pants.engine.platform import Platform
+from pants.util.enums import match
 
 
 # NB: We manually implement __hash__(), rather than using unsafe_hash=True, because PayloadField
@@ -38,7 +39,7 @@ def alias(cls):
 
   def as_shared_lib(self, platform):
     # TODO: check that the name conforms to some format in the constructor (e.g. no dots?).
-    return platform.match({
+    return match(platform, {
       Platform.darwin: f"lib{self.lib_name}.dylib",
       Platform.linux: f"lib{self.lib_name}.so",
     })