apache-ignite

1
# Licensed to the Apache Software Foundation (ASF) under one or more
2
# contributor license agreements.  See the NOTICE file distributed with
3
# this work for additional information regarding copyright ownership.
4
# The ASF licenses this file to You under the Apache License, Version 2.0
5
# (the "License"); you may not use this file except in compliance with
6
# the License.  You may obtain a copy of the License at
7
#
8
#    http://www.apache.org/licenses/LICENSE-2.0
9
#
10
# Unless required by applicable law or agreed to in writing, software
11
# distributed under the License is distributed on an "AS IS" BASIS,
12
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
# See the License for the specific language governing permissions and
14
# limitations under the License.
15

16
"""
17
This module contains various syntaxic sugar for enums
18
"""
19

20
from enum import IntEnum
21

22

23
def constructible(cls):
24
    """
25
    If decorate IntEnum subclass with this decorator, following is possible:
26

27
    @constructible
28
    class StopType(IntEnum):
29
        SIGTERM = 0
30
        SIGKILL = 1
31
        DISCONNECT = 2
32

33
    stop_type = 1
34

35
    with StopType.construct_from(stop_type) as sp:
36
        if sp is StopType.SIGKILL:
37
            print('got sigkill')
38

39
    stop_type = 'DISCONNECT'
40

41
    with StopType.construct_from(stop_type) as sp:
42
        if sp is StopType.DISCONNECT:
43
            print('got disconnect')
44

45
    """
46
    assert issubclass(cls, IntEnum)
47

48
    members = dict(cls.__members__.items())
49

50
    def construct_from(val):
51
        if isinstance(val, str):
52
            return members[val]
53
        return cls.__new__(cls, val)
54

55
    def __enter__(self):
56
        return self
57

58
    def __exit__(self, *args, **kwargs):
59
        pass
60

61
    cls.construct_from = construct_from
62
    cls.__enter__ = __enter__
63
    cls.__exit__ = __exit__
64

65
    return cls
66