Remove Python 3.7 support and fix ugly formatting in pyglove core

PiperOrigin-RevId: 597256649

vizier/_src/algorithms/designers/gp/yjt.py

@@ -16,7 +16,7 @@
 
 """Temporary file for finding the optimal Yeo-Johnson transformation."""
 
-
+from typing import Literal
 
 from absl import logging
 import numpy as np
@@ -32,7 +32,7 @@
 # dependency on sklearn.
 def optimal_transformation(
  data: np.ndarray,
- method: str = 'yeo-johnson',
+ method: Literal['yeo-johnson', 'box-cox'] = 'yeo-johnson',
  *,
  standardize: bool = True) -> tfb.AutoCompositeTensorBijector:
  """Returns the power transformation with optimal parameterization.

vizier/_src/benchmarks/experimenters/synthetic/simplekd.py

@@ -22,7 +22,7 @@
 parameter which make it harder to optimize.
 """
 
-, Sequence, Union
+from typing import Literal, Sequence, Union
 import attrs
 import numpy as np
 from vizier import pyvizier as vz
@@ -48,7 +48,7 @@ def _float_term(x_list: list[float]) -> float:
  return float_term
 
 
-SimpleKDCategory = str
+SimpleKDCategory = Literal['corner', 'center', 'mixed']
 
 
 def _categorical_term(x: str, best_category: SimpleKDCategory) -> float:

vizier/_src/pyglove/converters.py

@@ -443,7 +443,7 @@ def to_dna(self, trial: vz.Trial) -> pg.DNA:
  return dna
 
  def to_trial(self, dna: pg.DNA, *,
- fallback: str) -> vz.Trial:
+ fallback: Literal['raise_error', 'return_dummy']) -> vz.Trial:
  """Converts DNA to vizier Trial.
 
  Args:

vizier/_src/pyglove/core.py

@@ -56,6 +56,9 @@ class VizierTrial(pg.tuning.Trial):
  def __init__(
  self, converter: converters.VizierConverter, trial: vz.Trial, **kwargs
  ):
+ completed_time = (
+ int(trial.completion_time.timestamp()) if trial.completion_time else 0
+ )
  super().__init__(
  dna=pg.DNA(None),
  id=trial.id,
@@ -65,11 +68,7 @@ def __init__(
  ),
  status=_trial_status_legacy_value(trial.status),
  created_time=int(trial.creation_time.timestamp()),
- completed_time=int(
- (
- trial.completion_time or datetime.datetime.fromtimestamp(0)
- ).timestamp()
- ),
+ completed_time=completed_time,
  infeasible=trial.infeasible,
  **kwargs,
  )

vizier/_src/pyvizier/oss/metadata_util.py

@@ -16,7 +16,7 @@
 
 """Utility functions for handling vizier metadata."""
 
-from typing import Dict, Iterable, Optional, Tuple, Type, TypeVar, Union
+from typing import Dict, Iterable, Literal, Optional, Tuple, Type, TypeVar, Union
 from absl import logging
 
 from vizier._src.pyvizier.shared import common
@@ -51,7 +51,7 @@ def assign(
  key: str,
  ns: str,
  value: Union[str, any_pb2.Any, Message],
- mode: str = 'insert',
+ mode: Literal['insert_or_assign', 'insert_or_error', 'insert'] = 'insert',
 ) -> Tuple[key_value_pb2.KeyValue, bool]:
  """Insert and/or assign (key, value) to container.metadata.
 

vizier/_src/pyvizier/shared/parameter_iterators.py

@@ -18,7 +18,7 @@
 
 from typing import Iterator
 import copy
-from typing import Generator, Union
+from typing import Generator, Literal, Union
 
 from vizier._src.pyvizier.shared.parameter_config import ParameterConfig
 from vizier._src.pyvizier.shared.parameter_config import SearchSpace
@@ -48,7 +48,7 @@ def __init__(self,
  search_space: SearchSpace,
 
  *,
- traverse_order: str = 'dfs'):
+ traverse_order: Literal['dfs', 'bfs'] = 'dfs'):
  """Init.
 
  See the class pydoc for more details.

vizier/_src/pyvizier/shared/parameter_iterators_test.py

@@ -17,7 +17,7 @@
 """Tests for parameter iterators."""
 
 from typing import Sequence
-
+from typing import Literal
 
 from vizier import pyvizier as vz
 from vizier._src.pyvizier.shared import parameter_iterators as pi
@@ -39,7 +39,7 @@ class ParameterIteratorsTest(parameterized.TestCase):
  expected_order=[
  'model', 'num_layers', 'apply_preprocessing', 'preprocessor'
  ]))
- def test_e2e(self, traverse_order: str,
+ def test_e2e(self, traverse_order: Literal['bfs', 'dfs'],
  expected_order: Sequence[str]):
  valid_params = {
  'model': 'dnn',