Coverage for .nox/test-3-9/lib/python3.9/site-packages/nskit/common/configuration/mixins.py: 97%

29 statements  

« prev     ^ index     » next       coverage.py v7.4.2, created at 2024-02-25 17:38 +0000

1"""Configuration Mixins.""" 

2import inspect 

3from typing import Any, Callable 

4 

5from pydantic import BaseModel, model_serializer, SerializationInfo 

6from pydantic_settings import BaseSettings 

7 

8 

9class PropertyDumpMixin: 

10 """Mixin to dump properties when doing model_dump.""" 

11 

12 def model_post_init(self, __context: Any): 

13 """Make sure excluded properties is created.""" 

14 if not hasattr(self, '_excluded_properties') or self._excluded_properties is None: 

15 self._excluded_properties = [] 

16 super().model_post_init(__context) 

17 

18 @model_serializer(mode='wrap') 

19 def _property_dump(self, handler: Callable, info: SerializationInfo): 

20 value = handler(self) 

21 # Add properties if they are included/excluded 

22 value.update(self.__get_properties(info.include, info.exclude)) 

23 return value 

24 

25 def __get_defined_model_properties(self): 

26 

27 properties = inspect.getmembers(self.__class__, lambda o: isinstance(o, property)) 

28 # Based on object -> BaseModel -> BaseSettings 

29 standard_settings_properties = inspect.getmembers(BaseSettings, lambda o: isinstance(o, property)) 

30 standard_model_properties = inspect.getmembers(BaseModel, lambda o: isinstance(o, property)) 

31 property_names = [u[0] for u in properties if u not in standard_settings_properties+standard_model_properties and not u[0].startswith('_')] 

32 return property_names 

33 

34 def __get_properties(self, include=None, exclude=None): 

35 property_names = self.__get_defined_model_properties() 

36 included_properties = {} 

37 if not hasattr(self, '_excluded_properties') or self._excluded_properties is None: 

38 self._excluded_properties = [] 

39 for property_name in property_names: 

40 if (include and property_name in include) or \ 

41 (exclude and property_name not in exclude and property_name not in self._excluded_properties) or \ 

42 (include is None and exclude is None and property_name not in self._excluded_properties): 

43 included_properties[property_name] = getattr(self, property_name) 

44 return included_properties