Phase 2: Migrate to StrEnum (#4868)

Summary:
Pull Request resolved: #4868

Convert (str, Enum) and Enum patterns to Python 3.11's built-in StrEnum class.
This removes the need for the dual inheritance pattern and removes pyre-fixme
annotations that were needed to annotate enum values with str type.

Files changed:
- ax/utils/common/constants.py - Keys(str, Enum) -> Keys(StrEnum)
- ax/storage/utils.py - MetricIntent(enum.Enum) -> MetricIntent(StrEnum)
- ax/core/auxiliary.py - AuxiliaryExperimentPurpose(Enum) -> AuxiliaryExperimentPurpose(StrEnum)
- ax/utils/stats/model_fit_stats.py - ModelFitMetricDirection(Enum) -> ModelFitMetricDirection(StrEnum)
- ax/metrics/chemistry.py - ChemistryProblemType(Enum) -> ChemistryProblemType(StrEnum)
- ax/metrics/sklearn.py - SklearnModelType(Enum), SklearnDataset(Enum) -> StrEnum

Reviewed By: saitcakmak

Differential Revision: D91648882

fbshipit-source-id: 56eaaced06aafe8dc1743f6633562e6afcf3224f

Author: Balandat

ax/core/auxiliary.py

@@ -9,7 +9,7 @@
 
 from abc import ABC
 from dataclasses import dataclass
-from enum import Enum, unique
+from enum import StrEnum, unique
 from typing import TYPE_CHECKING
 
 from ax.core.data import Data
@@ -51,7 +51,7 @@ def _unique_id(self) -> str:
 
 
 @unique
-class AuxiliaryExperimentPurpose(Enum):
+class AuxiliaryExperimentPurpose(StrEnum):
     # BOPE Aux Experiment Usage pattern:
     # 1. Run the exploratory batch for the main / BO experiment.
     # 2. Use the BO experiment as the auxiliary experiment for the PE experiment

ax/metrics/chemistry.py

@@ -33,7 +33,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from enum import Enum
+from enum import StrEnum
 from functools import lru_cache
 from pathlib import Path
 from typing import Any
@@ -48,11 +48,9 @@
 from pyre_extensions import none_throws
 
 
-class ChemistryProblemType(Enum):
-    # pyre-fixme[35]: Target cannot be annotated.
-    SUZUKI: str = "suzuki"
-    # pyre-fixme[35]: Target cannot be annotated.
-    DIRECT_ARYLATION: str = "direct_arylation"
+class ChemistryProblemType(StrEnum):
+    SUZUKI = "suzuki"
+    DIRECT_ARYLATION = "direct_arylation"
 
 
 @dataclass(frozen=True)

ax/metrics/sklearn.py

@@ -9,7 +9,7 @@
 from __future__ import annotations
 
 from copy import deepcopy
-from enum import Enum
+from enum import StrEnum
 from functools import lru_cache
 from math import sqrt
 from typing import Any
@@ -28,20 +28,15 @@
 from sklearn.neural_network import MLPClassifier, MLPRegressor
 
 
-class SklearnModelType(Enum):
-    # pyre-fixme[35]: Target cannot be annotated.
-    RF: str = "rf"
-    # pyre-fixme[35]: Target cannot be annotated.
-    NN: str = "nn"
+class SklearnModelType(StrEnum):
+    RF = "rf"
+    NN = "nn"
 
 
-class SklearnDataset(Enum):
-    # pyre-fixme[35]: Target cannot be annotated.
-    DIGITS: str = "digits"
-    # pyre-fixme[35]: Target cannot be annotated.
-    BOSTON: str = "boston"
-    # pyre-fixme[35]: Target cannot be annotated.
-    CANCER: str = "cancer"
+class SklearnDataset(StrEnum):
+    DIGITS = "digits"
+    BOSTON = "boston"
+    CANCER = "cancer"
 
 
 @lru_cache(maxsize=8)

ax/storage/utils.py

@@ -9,6 +9,7 @@
 import enum
 from collections import OrderedDict
 from collections.abc import Mapping
+from enum import StrEnum
 from hashlib import md5
 
 import pandas as pd
@@ -25,7 +26,7 @@ class DomainType(enum.Enum):
     DERIVED = 4
 
 
-class MetricIntent(enum.Enum):
+class MetricIntent(StrEnum):
     """Class for enumerating metric use types."""
 
     OBJECTIVE = "objective"

ax/utils/common/constants.py

@@ -6,7 +6,7 @@
 
 # pyre-strict
 
-from enum import Enum, unique
+from enum import StrEnum, unique
 
 # ------------------------- Miscellaneous -------------------------
 
@@ -36,7 +36,7 @@
 
 
 @unique
-class Keys(str, Enum):
+class Keys(StrEnum):
     """Enum of reserved keys in options dicts etc, alphabetized.
 
 

ax/utils/stats/model_fit_stats.py

@@ -6,7 +6,7 @@
 # pyre-strict
 
 from collections.abc import Mapping
-from enum import Enum
+from enum import StrEnum
 from logging import Logger
 from typing import Protocol
 
@@ -33,7 +33,7 @@
 KENDALL_TAU_RANK_CORRELATION = "Kendall tau rank correlation"
 
 
-class ModelFitMetricDirection(Enum):
+class ModelFitMetricDirection(StrEnum):
     """Model fit metric directions."""
 
     MINIMIZE = "minimize"