"""Policy Supporter used for the OSS Vizier service.

The Policy can use these methods to communicate with Vizier.
"""
import datetime
from typing import Iterable, List, Optional, Union

from vizier import pythia
from vizier import pyvizier as vz
from vizier.service import pythia_util
from vizier.service import pyvizier
from vizier.service import vizier_service_pb2
from vizier.service import vizier_service_pb2_grpc

VizierService = Union[vizier_service_pb2_grpc.VizierServiceStub,
                      vizier_service_pb2_grpc.VizierServiceServicer]


# TODO: Consider replacing protos with Pyvizier clients.py.
class ServicePolicySupporter(pythia.PolicySupporter):
  """Service version of the PolicySupporter."""

  def __init__(self, study_guid: str, vizier_service: VizierService):
    """Initalization.

    Args:
      study_guid: A default study_name; the name of this study.
      vizier_service: Vizier Service, in the form of a GRPC stub or actual
        class.
    """
    self._study_guid = study_guid
    self._vizier_service = vizier_service

  def GetStudyConfig(self,
                     study_guid: Optional[str] = None) -> vz.ProblemStatement:
    if study_guid is None:
      study_guid = self._study_guid
    request = vizier_service_pb2.GetStudyRequest(name=study_guid)
    study = self._vizier_service.GetStudy(request)
    return pyvizier.StudyConfig.from_proto(study.study_spec).to_problem()

  # TODO: Use TrialsFilter instead.
  def GetTrials(
      self,
      *,
      study_guid: Optional[str] = None,  # Equivalent to Study.name in OSS.
      trial_ids: Optional[Iterable[int]] = None,
      min_trial_id: Optional[int] = None,
      max_trial_id: Optional[int] = None,
      status_matches: Optional[vz.TrialStatus] = None,
      include_intermediate_measurements: bool = True) -> List[vz.Trial]:
    """Requests Trials from Vizier."""

    if study_guid is None:
      study_guid = self._study_guid
    request = vizier_service_pb2.ListTrialsRequest(parent=study_guid)
    trials = self._vizier_service.ListTrials(request).trials
    all_pytrials = pyvizier.TrialConverter.from_protos(trials)

    if trial_ids is not None:
      # We're going to do repeated comparisons in _trial_filter(), so this
      # needs to be more real than an Iterator.
      trial_ids = frozenset(trial_ids)

    def _trial_filter(trial) -> bool:
      if trial_ids is not None and trial.id not in trial_ids:
        return False
      else:
        if min_trial_id is not None:
          if trial.id < min_trial_id:
            return False
        if max_trial_id is not None:
          if trial.id > max_trial_id:
            return False
        if status_matches:
          if trial.status != status_matches:
            return False
      return True

    filtered_pytrials = [t for t in all_pytrials if _trial_filter(t)]

    # Doesn't affect datastore when measurements are deleted.
    if not include_intermediate_measurements:
      for filtered_pytrial in filtered_pytrials:
        filtered_pytrial.measurements = []

    return filtered_pytrials

  def CheckCancelled(self, note: Optional[str] = None) -> None:
    """Throws a CancelComputeError on timeout or if Vizier cancels."""
    pass  # Do nothing since it's one single process.

  def TimeRemaining(self) -> datetime.timedelta:
    """The time remaining to compute a result."""
    return datetime.timedelta.max  # RPCs don't have timeouts in OSS.

  def SendMetadata(self, delta: vz.MetadataDelta) -> None:
    """Updates the metadata."""
    self.CheckCancelled('UpdateMetadata entry')
    request = pyvizier.metadata_util.to_request_proto(self._study_guid, delta)
    waiter = pythia_util.ResponseWaiter()
    # In a real server, this happens in another thread:
    response = self._vizier_service.UpdateMetadata(request)
    waiter.Report(response)
    # Wait for Vizier to reply with a UpdateMetadataResponse packet.
    waiter.WaitForResponse()
    self.CheckCancelled('UpdateMetadata exit')

@@ -1084,150 +1084,49 @@
 elf,
-%0A                     study_guid: Optional%5Bstr%5D = None) -%3E vz.ProblemStatement:%0A    if study_guid is None:%0A      study_guid = self._study_guid
+ study_guid: str) -%3E vz.ProblemStatement:
 %0A   
@@ -1328,32 +1328,44 @@
 DO: 
-Use TrialsF
+Support f
 ilter
+s
  in
-stead
+ ListTrialsRequest
 .%0A  
@@ -1748,33 +1748,48 @@
  %22%22%22
-Requests Trials from Vizi
+Fetch all trials and then apply the filt
 er.%22
@@ -2058,329 +2058,61 @@
     
-if trial_ids is not None:%0A      # We're going to do repeated comparisons in _trial_filter(), so this%0A      # needs to be more real than an Iterator.%0A      trial_ids = frozenset(trial_ids)%0A%0A    def _trial_filter(trial) -%3E bool:%0A      if trial_ids is not None and trial.id not in trial_ids:%0A        return False%0A      else:
+trial_filter = vz.TrialFilter(%0A        ids=trial_ids,
 %0A   
@@ -2120,61 +2120,15 @@
     
-if 
 min_
-trial_id is not None:%0A          if trial.id %3C 
+id=
 min_
@@ -2131,35 +2131,34 @@
 min_trial_id
-:%0A 
+,%0A
         
    return Fa
@@ -2149,85 +2149,15 @@
     
-   return False%0A        if max_trial_id is not None:%0A          if trial.id %3E 
+max_id=
 max_
@@ -2164,27 +2164,26 @@
 trial_id
-:%0A 
+,%0A
         
    retur
@@ -2178,139 +2178,60 @@
     
-   return False%0A        if status_matches:%0A          if trial.status != status_matches:%0A            return False%0A      return True%0A
+status=%5Bstatus_matches%5D if status_matches else None)
 %0A   
@@ -2279,17 +2279,16 @@
 ials if 
-_
 trial_fi

Use TrialFilter in ServicePolicySupporter and follow exact PolicySupporter base class API.