Fix Enum.Value identity under pickling.

pex/atomic_directory.py

@@ -181,6 +181,9 @@ class Value(Enum.Value):
     POSIX = Value("posix")
 
 
+FileLockStyle.seal()
+
+
 def _is_bsd_lock(lock_style=None):
     # type: (Optional[FileLockStyle.Value]) -> bool
 

pex/bin/pex.py

@@ -589,6 +589,9 @@ class Value(Enum.Value):
     VERBOSE = Value("verbose")
 
 
+Seed.seal()
+
+
 class HandleSeedAction(Action):
     def __init__(self, *args, **kwargs):
         kwargs["nargs"] = "?"

pex/cache/dirs.py

@@ -204,6 +204,8 @@ def iter_transitive_dependents(self):
     )
 
 
+CacheDir.seal()
+
 if TYPE_CHECKING:
     _AtomicCacheDir = TypeVar("_AtomicCacheDir", bound="AtomicCacheDir")
 

pex/cli/commands/cache/bytes.py

@@ -41,6 +41,9 @@ def render(self, total_bytes):
     PB = Value("PB", 1000 * TB.multiple)
 
 
+ByteUnits.seal()
+
+
 @attr.s(frozen=True)
 class ByteAmount(object):
     @classmethod

pex/cli/commands/lock.py

@@ -89,6 +89,9 @@ class Value(Enum.Value):
     ERROR = Value("error")
 
 
+FingerprintMismatch.seal()
+
+
 class ExportFormat(Enum["ExportFormat.Value"]):
     class Value(Enum.Value):
         pass
@@ -98,6 +101,9 @@ class Value(Enum.Value):
     PEP_665 = Value("pep-665")
 
 
+ExportFormat.seal()
+
+
 class ExportSortBy(Enum["ExportSortBy.Value"]):
     class Value(Enum.Value):
         pass
@@ -106,6 +112,9 @@ class Value(Enum.Value):
     PROJECT_NAME = Value("project-name")
 
 
+ExportSortBy.seal()
+
+
 class DryRunStyle(Enum["DryRunStyle.Value"]):
     class Value(Enum.Value):
         pass
@@ -114,6 +123,9 @@ class Value(Enum.Value):
     CHECK = Value("check")
 
 
+DryRunStyle.seal()
+
+
 class HandleDryRunAction(Action):
     def __init__(self, *args, **kwargs):
         kwargs["nargs"] = "?"

pex/cli/commands/venv.py

@@ -49,6 +49,9 @@ class Value(Enum.Value):
     FLAT_ZIPPED = Value("flat-zipped")
 
 
+InstallLayout.seal()
+
+
 class Venv(OutputMixin, JsonMixin, BuildTimeCommand):
     @classmethod
     def _add_inspect_arguments(cls, parser):

pex/common.py

@@ -855,6 +855,9 @@ class Value(Enum.Value):
     SYMLINK = Value("symlink")
 
 
+CopyMode.seal()
+
+
 def iter_copytree(
     src,  # type: Text
     dst,  # type: Text

pex/dist_metadata.py

@@ -201,6 +201,9 @@ def load_metadata(
     PKG_INFO = Value("PKG-INFO")
 
 
+MetadataType.seal()
+
+
 @attr.s(frozen=True)
 class MetadataKey(object):
     metadata_type = attr.ib()  # type: MetadataType.Value
@@ -963,6 +966,9 @@ def of(cls, location):
         return cls.SDIST
 
 
+DistributionType.seal()
+
+
 @attr.s(frozen=True)
 class Distribution(object):
     @staticmethod

pex/enum.py

@@ -3,20 +3,34 @@
 
 from __future__ import absolute_import, print_function
 
+import sys
 import weakref
 from collections import defaultdict
 from functools import total_ordering
 
 from _weakref import ReferenceType
 
+from pex.exceptions import production_assert
 from pex.typing import TYPE_CHECKING, Generic, cast
 
 if TYPE_CHECKING:
-    from typing import Any, DefaultDict, List, Optional, Tuple, Type, TypeVar
+    from typing import Any, DefaultDict, Iterator, List, Optional, Tuple, Type, TypeVar
 
     _V = TypeVar("_V", bound="Enum.Value")
 
 
+def _get_or_create(
+    module,  # type: str
+    enum_type,  # type: str
+    enum_value_type,  # type: str
+    enum_value_value,  # type: str
+):
+    # type: (...) -> Enum.Value
+    enum_class = getattr(sys.modules[module], enum_type)
+    enum_value_class = getattr(enum_class, enum_value_type)
+    return cast("Enum.Value", enum_value_class._get_or_create(enum_value_value))
+
+
 class Enum(Generic["_V"]):
     @total_ordering
     class Value(object):
@@ -26,11 +40,37 @@ class Value(object):
 
         @classmethod
         def _iter_values(cls):
+            # type: () -> Iterator[Enum.Value]
             for ref in cls._values_by_type[cls]:
                 value = ref()
                 if value:
                     yield value
 
+        @classmethod
+        def _get_or_create(cls, value):
+            # type: (str) -> Enum.Value
+            for existing_value in cls._iter_values():
+                if existing_value.value == value:
+                    return existing_value
+            return cls(value)
+
+        def __reduce__(self):
+            if sys.version_info[0] >= 3:
+                return self._get_or_create, (self.value,)
+
+            # N.B.: Python 2.7 does not handle pickling nested classes; so we go through some
+            # hoops here and in `Enum.seal`.
+            module = self.__module__
+            enum_type = getattr(self, "_enum_type", None)
+            production_assert(
+                isinstance(enum_type, str),
+                "The Enum subclass in the {module} module containing value {self} was not "
+                "`seal`ed.",
+                module=module,
+                self=self,
+            )
+            return _get_or_create, (module, enum_type, type(self).__name__, self.value)
+
         def __init__(self, value):
             # type: (str) -> None
             values = Enum.Value._values_by_type[type(self)]
@@ -78,6 +118,46 @@ def __le__(self, other):
                 raise self._create_type_error(other)
             return self is other or self < other
 
+    @classmethod
+    def seal(cls):
+        if sys.version_info[0] >= 3:
+            return
+
+        # N.B.: Python 2.7 does not handle pickling nested classes; so we go through some
+        # hoops here and in `Enum.Value.__reduce__`.
+
+        enum_type_name, _, enum_value_type_name = cls.type_var.partition(".")
+        if enum_value_type_name:
+            production_assert(
+                cls.__name__ == enum_type_name,
+                "Expected Enum subclass {cls} to have a type parameter of the form `{name}.Value` "
+                "where `Value` is a subclass of `Enum.Value`. Instead found: {type_var}",
+                cls=cls,
+                name=cls.__name__,
+                type_var=cls.type_var,
+            )
+            enum_value_type = getattr(cls, enum_value_type_name, None)
+        else:
+            enum_value_type = getattr(sys.modules[cls.__module__], enum_type_name, None)
+
+        production_assert(
+            enum_type_name is not None,
+            "Failed to find Enum.Value type {type_var} for Enum {cls} in module {module}",
+            type_var=cls.type_var,
+            cls=cls,
+            module=cls.__module__,
+        )
+        production_assert(
+            issubclass(enum_value_type, Enum.Value),
+            "Expected Enum subclass {cls} to have a type parameter that is a subclass of "
+            "`Enum.Value`. Instead found {type_var} was of type: {enum_value_type}",
+            cls=cls,
+            name=cls.__name__,
+            type_var=cls.type_var,
+            enum_value_type=enum_value_type,
+        )
+        setattr(enum_value_type, "_enum_type", cls.__name__)
+
     _values = None  # type: Optional[Tuple[_V, ...]]
 
     @classmethod

pex/inherit_path.py

@@ -33,3 +33,6 @@ def for_value(cls, value):
                     value, type(value)
                 )
             )
+
+
+InheritPath.seal()

pex/interpreter_constraints.py

@@ -328,6 +328,8 @@ class Value(Enum.Value):
     EOL = Value("eol")
 
 
+Lifecycle.seal()
+
 # This value is based off of:
 # 1. Past releases: https://www.python.org/downloads/ where the max patch level was achieved by
 #    2.7.18.

pex/layout.py

@@ -75,6 +75,9 @@ def identify_original(cls, pex):
         return cls.Value.try_load(pex) or Layout.LOOSE
 
 
+Layout.seal()
+
+
 class _Layout(object):
     def __init__(
         self,

pex/pep_427.py

@@ -53,6 +53,9 @@ class Value(Enum.Value):
     WHEEL_FILE = Value(".whl file")
 
 
+InstallableType.seal()
+
+
 @attr.s(frozen=True)
 class InstallPaths(object):
 

pex/pex_builder.py

@@ -112,6 +112,9 @@ def perform_check(
     ERROR = Value("error")
 
 
+Check.seal()
+
+
 class PEXBuilder(object):
     """Helper for building PEX environments."""
 

pex/pip/version.py

@@ -320,3 +320,6 @@ def values(cls):
     VENDORED = v20_3_4_patched
     LATEST = LatestPipVersion()
     DEFAULT = DefaultPipVersion(preferred=(VENDORED, v23_2, v24_1))
+
+
+PipVersion.seal()

pex/requirements.py

@@ -164,6 +164,9 @@ class Value(Enum.Value):
     Subversion = Value("svn")
 
 
+VCS.seal()
+
+
 @attr.s(frozen=True)
 class VCSRequirement(object):
     """A requirement realized by building a distribution from sources retrieved from a VCS."""
@@ -278,6 +281,9 @@ class Value(Enum.Value):
     HTTPS = Value("https")
 
 
+ArchiveScheme.seal()
+
+
 @attr.s(frozen=True)
 class VCSScheme(object):
     vcs = attr.ib()  # type: VCS.Value

pex/resolve/locked_resolve.py

@@ -67,6 +67,9 @@ class Value(Enum.Value):
     UNIVERSAL = Value("universal")
 
 
+LockStyle.seal()
+
+
 class TargetSystem(Enum["TargetSystem.Value"]):
     class Value(Enum.Value):
         pass
@@ -76,6 +79,9 @@ class Value(Enum.Value):
     WINDOWS = Value("windows")
 
 
+TargetSystem.seal()
+
+
 @attr.s(frozen=True)
 class LockConfiguration(object):
     style = attr.ib()  # type: LockStyle.Value

pex/resolve/resolver_configuration.py

@@ -57,6 +57,9 @@ def default(cls, pip_version=None):
     PIP_2020 = Value("pip-2020-resolver")
 
 
+ResolverVersion.seal()
+
+
 @attr.s(frozen=True)
 class ReposConfiguration(object):
     @classmethod

pex/scie/model.py

@@ -35,6 +35,9 @@ class Value(Enum.Value):
     EAGER = Value("eager")
 
 
+ScieStyle.seal()
+
+
 class PlatformNamingStyle(Enum["PlatformNamingStyle.Value"]):
     class Value(Enum.Value):
         pass
@@ -44,6 +47,9 @@ class Value(Enum.Value):
     FILE_SUFFIX = Value("platform-file-suffix")
 
 
+PlatformNamingStyle.seal()
+
+
 @attr.s(frozen=True)
 class ConsoleScript(object):
     name = attr.ib()  # type: str
@@ -295,6 +301,9 @@ def parse(cls, value):
         return cls.CURRENT if "current" == value else cls.for_value(value)
 
 
+SciePlatform.seal()
+
+
 class Provider(Enum["Provider.Value"]):
     class Value(Enum.Value):
         pass
@@ -303,6 +312,9 @@ class Value(Enum.Value):
     PyPy = Value("PyPy")
 
 
+Provider.seal()
+
+
 @attr.s(frozen=True)
 class InterpreterDistribution(object):
     provider = attr.ib()  # type: Provider.Value

pex/tools/commands/venv.py

@@ -39,6 +39,9 @@ class Value(Enum.Value):
     PEX_AND_PEX_ROOT = Value("all")
 
 
+RemoveScope.seal()
+
+
 @attr.s(frozen=True)
 class InstallScopeState(object):
     @classmethod

pex/venv/bin_path.py

@@ -13,3 +13,6 @@ class Value(Enum.Value):
     FALSE = Value("false")
     PREPEND = Value("prepend")
     APPEND = Value("append")
+
+
+BinPath.seal()