FEDOT

1
from typing import Optional, Union
2

3
from fedot.core.operations.evaluation.operation_implementations.implementation_interfaces import \
4
    DataOperationImplementation, ModelImplementation
5
from fedot.core.utils import RandomStateHandler
6

7

8
class ImplementationRandomStateHandler(RandomStateHandler):
9
    def __init__(self, seed: Optional[int] = None,
10
                 implementation: Union[DataOperationImplementation, ModelImplementation] = None):
11
        super().__init__(seed)
12
        self.implementation = implementation
13
        self.implementation_old_random_state = None
14

15
    def __enter__(self):
16
        self.implementation_old_random_state = _set_operation_random_seed(self.implementation, self._seed)
17
        return super().__enter__()
18

19
    def __exit__(self, exc_type, exc_value, exc_traceback):
20
        _set_operation_random_seed(self.implementation, self.implementation_old_random_state)
21
        return super().__exit__(exc_type, exc_value, exc_traceback)
22

23

24
def _set_operation_random_seed(operation: Union[DataOperationImplementation, ModelImplementation],
25
                               seed: Optional[int]):
26
    old_random_state = None
27
    if hasattr(operation, 'random_state'):
28
        old_random_state = getattr(operation, 'random_state')
29
        setattr(operation, 'random_state', seed)
30

31
    elif hasattr(operation, 'seed'):
32
        old_random_state = getattr(operation, 'seed')
33
        setattr(operation, 'seed', seed)
34

35
    return old_random_state
36