Coverage for src/mlopus/mlflow/api/entity.py: 100%

1from abc import abstractmethod 

2from typing import Mapping 

3 

4from mlopus.utils import pydantic 

5from . import contract 

6from .common import schema 

7 

8 

9class EntityApi(schema.BaseEntity): 

10 """Self-updating entity with API handle.""" 

11 

12 api: contract.MlflowApiContract = pydantic.Field(default=None, exclude=True, repr=False) 

13 

14 def using(self, api: contract.MlflowApiContract) -> "EntityApi": 

15 """Switch API handle.""" 

16 self.api = api 

17 return self 

18 

19 def update(self) -> "EntityApi": 

20 """Use API to get latest data for this entity and update this entity in place.""" 

21 return self._use_values_from(self._get_latest_data()) 

22 

23 def _use_values_from(self, other: schema.BaseEntity) -> "EntityApi": 

24 """Update this entity in place with values from the other entity. Return a reference to this entity.""" 

25 _ = [setattr(self, k, v) for k, v in other.items()] 

26 return self 

27 

28 @abstractmethod 

29 def _get_latest_data(self) -> schema.BaseEntity: 

30 """Use API to get latest data for this entity.""" 

31 

32 @abstractmethod 

33 def set_tags(self, tags: Mapping) -> "EntityApi": 

34 """Set tags on this entity."""