8000 simplified enums NodeStatus and IsolationLevel · postgrespro/testgres@cf03ee2 · GitHub
[go: up one dir, main page]

Skip to content

Commit cf03ee2

Browse files
committed
simplified enums NodeStatus and IsolationLevel
1 parent 5878552 commit cf03ee2

File tree

2 files changed

+13
-28
lines changed

2 files changed

+13
-28
lines changed

testgres/connection.py

Lines changed: 10 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,10 @@ class IsolationLevel(Enum):
2727
Transaction isolation level for NodeConnection
2828
"""
2929

30-
ReadUncommitted, ReadCommitted, RepeatableRead, Serializable = range(4)
30+
ReadUncommitted = 'read uncommitted'
31+
ReadCommitted = 'read committed'
32+
RepeatableRead = 'repeatable read'
33+
Serializable = 'serializable'
3134

3235

3336
class NodeConnection(object):
@@ -71,39 +74,21 @@ def __exit__(self, type, value, traceback):
7174
self.close()
7275

7376
def begin(self, isolation_level=IsolationLevel.ReadCommitted):
74-
# yapf: disable
75-
levels = [
76-
'read uncommitted',
77-
'read committed',
78-
'repeatable read',
79-
'serializable'
80-
]
81-
82-
# Check if level is an IsolationLevel
83-
if (isinstance(isolation_level, IsolationLevel)):
84-
85-
# Get index of isolation level
86-
level_idx = isolation_level.value
87-
assert level_idx in range(4)
88-
89-
# Replace isolation level with its name
90-
isolation_level = levels[level_idx]
91-
92-
else:
77+
# Check if level isn't an IsolationLevel
78+
if not isinstance(isolation_level, IsolationLevel):
9379
# Get name of isolation level
9480
level_str = str(isolation_level).lower()
9581

9682
# Validate level string
97-
if level_str not in levels:
83+
try:
84+
isolation_level = IsolationLevel(level_str)
85+
except ValueError:
9886
error = 'Invalid isolation level "{}"'
9987
raise QueryException(error.format(level_str))
10088

101-
# Replace isolation level with its name
102-
isolation_level = level_str
103-
10489
# Set isolation level
10590
cmd = 'SET TRANSACTION ISOLATION LEVEL {}'
106-
self.cursor.execute(cmd.format(isolation_level))
91+
self.cursor.execute(cmd.format(isolation_level.value))
10792

10893
return self
10994

testgres/node.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88
import tempfile
99
import time
1010

11-
from enum import Enum
11+
from enum import IntEnum
1212
from six import raise_from
1313

1414
from .cache import cached_initdb
@@ -53,7 +53,7 @@
5353
positional_args_hack
5454

5555

56-
class NodeStatus(Enum):
56+
class NodeStatus(IntEnum):
5757
"""
5858
Status of a PostgresNode
5959
"""
@@ -62,7 +62,7 @@ class NodeStatus(Enum):
6262

6363
# for Python 3.x
6464
def __bool__(self):
65-
return self.value == NodeStatus.Running.value
65+
return self == NodeStatus.Running
6666

6767
# for Python 2.x
6868
__nonzero__ = __bool__

0 commit comments

Comments
 (0)
0