[Author Prev][Author Next][Thread Prev][Thread Next][Author Index][Thread Index]

[tor-commits] [stem/master] Bundle enums with the numbers they map to



commit ab2dfc63394d0caa06a07ce14acc1fd0c784974e
Author: Damian Johnson <atagar@xxxxxxxxxxxxxx>
Date:   Tue Jan 23 11:57:56 2018 -0800

    Bundle enums with the numbers they map to
    
    Putting this all in one place both makes usage nicer, as well as less error
    prone for adding new values.
---
 stem/client/__init__.py | 162 +++++++++++++++++++++++++-----------------------
 stem/client/cell.py     |  57 +----------------
 2 files changed, 86 insertions(+), 133 deletions(-)

diff --git a/stem/client/__init__.py b/stem/client/__init__.py
index fcd0a238..a247c7ba 100644
--- a/stem/client/__init__.py
+++ b/stem/client/__init__.py
@@ -120,56 +120,93 @@ __all__ = [
   'cell',
 ]
 
-AddrType = stem.util.enum.UppercaseEnum(
-  'HOSTNAME',
-  'IPv4',
-  'IPv6',
-  'ERROR_TRANSIENT',
-  'ERROR_PERMANENT',
-  'UNKNOWN',
+
+class _IntegerEnum(stem.util.enum.Enum):
+  """
+  Integer backed enumeration. Enumerations of this type always have an implicit
+  **UNKNOWN** value for integer values that lack a mapping.
+  """
+
+  def __init__(self, *args):
+    self._enum_to_int = {}
+    self._int_to_enum = {}
+    parent_args = []
+
+    for entry in args:
+      if len(entry) == 2:
+        enum, int_val = entry
+        str_val = enum
+      elif len(entry) == 3:
+        enum, str_val, int_val = entry
+      else:
+        raise ValueError('IntegerEnums can only be constructed with two or three value tuples: %s' % repr(entry))
+
+      self._enum_to_int[enum] = int_val
+      self._int_to_enum[int_val] = enum
+      parent_args.append((enum, str_val))
+
+    parent_args.append(('UNKNOWN', 'UNKNOWN'))
+    super(_IntegerEnum, self).__init__(*parent_args)
+
+  def get(self, val):
+    """
+    Privides the (enum, int_value) tuple for a given value.
+    """
+
+    if isinstance(val, int):
+      return self._int_to_enum.get(val, self.UNKNOWN), val
+    elif val in self:
+      return val, self._enum_to_int.get(val, val)
+    else:
+      raise ValueError('Invalid %s type: %s' % (self.__name__, val))
+
+
+AddrType = _IntegerEnum(
+  ('HOSTNAME', 0),
+  ('IPv4', 4),
+  ('IPv6', 6),
+  ('ERROR_TRANSIENT', 16),
+  ('ERROR_PERMANENT', 17),
 )
 
-RelayCommand = stem.util.enum.Enum(
-  ('BEGIN', 'RELAY_BEGIN'),
-  ('DATA', 'RELAY_DATA'),
-  ('END', 'RELAY_END'),
-  ('CONNECTED', 'RELAY_CONNECTED'),
-  ('SENDME', 'RELAY_SENDME'),
-  ('EXTEND', 'RELAY_EXTEND'),
-  ('EXTENDED', 'RELAY_EXTENDED'),
-  ('TRUNCATE', 'RELAY_TRUNCATE'),
-  ('TRUNCATED', 'RELAY_TRUNCATED'),
-  ('DROP', 'RELAY_DROP'),
-  ('RESOLVE', 'RELAY_RESOLVE'),
-  ('RESOLVED', 'RELAY_RESOLVED'),
-  ('BEGIN_DIR', 'RELAY_BEGIN_DIR'),
-  ('EXTEND2', 'RELAY_EXTEND2'),
-  ('EXTENDED2', 'RELAY_EXTENDED2'),
-  ('UNKNOWN', 'UNKNOWN'),
+RelayCommand = _IntegerEnum(
+  ('BEGIN', 'RELAY_BEGIN', 1),
+  ('DATA', 'RELAY_DATA', 2),
+  ('END', 'RELAY_END', 3),
+  ('CONNECTED', 'RELAY_CONNECTED', 4),
+  ('SENDME', 'RELAY_SENDME', 5),
+  ('EXTEND', 'RELAY_EXTEND', 6),
+  ('EXTENDED', 'RELAY_EXTENDED', 7),
+  ('TRUNCATE', 'RELAY_TRUNCATE', 8),
+  ('TRUNCATED', 'RELAY_TRUNCATED', 9),
+  ('DROP', 'RELAY_DROP', 10),
+  ('RESOLVE', 'RELAY_RESOLVE', 11),
+  ('RESOLVED', 'RELAY_RESOLVED', 12),
+  ('BEGIN_DIR', 'RELAY_BEGIN_DIR', 13),
+  ('EXTEND2', 'RELAY_EXTEND2', 14),
+  ('EXTENDED2', 'RELAY_EXTENDED2', 15),
 )
 
-CertType = stem.util.enum.UppercaseEnum(
-  'LINK',
-  'IDENTITY',
-  'AUTHENTICATE',
-  'UNKNOWN',
+CertType = _IntegerEnum(
+  ('LINK', 1),
+  ('IDENTITY', 2),
+  ('AUTHENTICATE', 3),
 )
 
-CloseReason = stem.util.enum.UppercaseEnum(
-  'NONE',
-  'PROTOCOL',
-  'INTERNAL',
-  'REQUESTED',
-  'HIBERNATING',
-  'RESOURCELIMIT',
-  'CONNECTFAILED',
-  'OR_IDENTITY',
-  'OR_CONN_CLOSED',
-  'FINISHED',
-  'TIMEOUT',
-  'DESTROYED',
-  'NOSUCHSERVICE',
-  'UNKNOWN',
+CloseReason = _IntegerEnum(
+  ('NONE', 0),
+  ('PROTOCOL', 1),
+  ('INTERNAL', 2),
+  ('REQUESTED', 3),
+  ('HIBERNATING', 4),
+  ('RESOURCELIMIT', 5),
+  ('CONNECTFAILED', 6),
+  ('OR_IDENTITY', 7),
+  ('OR_CONN_CLOSED', 8),
+  ('FINISHED', 9),
+  ('TIMEOUT', 10),
+  ('DESTROYED', 11),
+  ('NOSUCHSERVICE', 12),
 )
 
 
@@ -298,25 +335,8 @@ class Address(Field):
   :var bytes value_bin: encoded address value
   """
 
-  TYPE_FOR_INT = {
-    0: AddrType.HOSTNAME,
-    4: AddrType.IPv4,
-    6: AddrType.IPv6,
-    16: AddrType.ERROR_TRANSIENT,
-    17: AddrType.ERROR_PERMANENT,
-  }
-
-  INT_FOR_TYPE = dict((v, k) for k, v in TYPE_FOR_INT.items())
-
   def __init__(self, addr_type, value):
-    if isinstance(addr_type, int):
-      self.type = Address.TYPE_FOR_INT.get(addr_type, AddrType.UNKNOWN)
-      self.type_int = addr_type
-    elif addr_type in AddrType:
-      self.type = addr_type
-      self.type_int = Address.INT_FOR_TYPE.get(addr_type, -1)
-    else:
-      raise ValueError('Invalid address type: %s' % addr_type)
+    self.type, self.type_int = AddrType.get(addr_type)
 
     if self.type == AddrType.IPv4:
       if stem.util.connection.is_valid_ipv4_address(value):
@@ -383,24 +403,8 @@ class Certificate(Field):
   :var bytes value: certificate value
   """
 
-  TYPE_FOR_INT = {
-    1: CertType.LINK,
-    2: CertType.IDENTITY,
-    3: CertType.AUTHENTICATE,
-  }
-
-  INT_FOR_TYPE = dict((v, k) for k, v in TYPE_FOR_INT.items())
-
   def __init__(self, cert_type, value):
-    if isinstance(cert_type, int):
-      self.type = Certificate.TYPE_FOR_INT.get(cert_type, CertType.UNKNOWN)
-      self.type_int = cert_type
-    elif cert_type in CertType:
-      self.type = cert_type
-      self.type_int = Certificate.INT_FOR_TYPE.get(cert_type, -1)
-    else:
-      raise ValueError('Invalid certificate type: %s' % cert_type)
-
+    self.type, self.type_int = CertType.get(cert_type)
     self.value = value
 
   def pack(self):
diff --git a/stem/client/cell.py b/stem/client/cell.py
index 0f292c10..9813a85e 100644
--- a/stem/client/cell.py
+++ b/stem/client/cell.py
@@ -44,7 +44,7 @@ import random
 import sys
 
 from stem import UNDEFINED
-from stem.client import ZERO, Address, Certificate, CloseReason, RelayCommand, Size, split
+from stem.client import ZERO, Address, Certificate, CloseReason, Size, split
 from stem.util import _hash_attr, datetime_to_unix
 
 FIXED_PAYLOAD_LEN = 509
@@ -265,26 +265,6 @@ class RelayCell(CircuitCell):
   VALUE = 3
   IS_FIXED_SIZE = True
 
-  COMMAND_FOR_INT = {
-    1: RelayCommand.BEGIN,
-    2: RelayCommand.DATA,
-    3: RelayCommand.END,
-    4: RelayCommand.CONNECTED,
-    5: RelayCommand.SENDME,
-    6: RelayCommand.EXTEND,
-    7: RelayCommand.EXTENDED,
-    8: RelayCommand.TRUNCATE,
-    9: RelayCommand.TRUNCATED,
-    10: RelayCommand.DROP,
-    11: RelayCommand.RESOLVE,
-    12: RelayCommand.RESOLVED,
-    13: RelayCommand.BEGIN_DIR,
-    14: RelayCommand.EXTEND2,
-    15: RelayCommand.EXTENDED2,
-  }
-
-  INT_FOR_COMMANDS = dict((v, k) for k, v in COMMAND_FOR_INT.items())
-
 
 class DestroyCell(CircuitCell):
   """
@@ -298,35 +278,9 @@ class DestroyCell(CircuitCell):
   VALUE = 4
   IS_FIXED_SIZE = True
 
-  REASON_FOR_INT = {
-    0: CloseReason.NONE,
-    1: CloseReason.PROTOCOL,
-    2: CloseReason.INTERNAL,
-    3: CloseReason.REQUESTED,
-    4: CloseReason.HIBERNATING,
-    5: CloseReason.RESOURCELIMIT,
-    6: CloseReason.CONNECTFAILED,
-    7: CloseReason.OR_IDENTITY,
-    8: CloseReason.OR_CONN_CLOSED,
-    9: CloseReason.FINISHED,
-    10: CloseReason.TIMEOUT,
-    11: CloseReason.DESTROYED,
-    12: CloseReason.NOSUCHSERVICE,
-  }
-
-  INT_FOR_REASON = dict((v, k) for k, v in REASON_FOR_INT.items())
-
   def __init__(self, circ_id, reason):
     super(DestroyCell, self).__init__(circ_id)
-
-    if isinstance(reason, int):
-      self.reason = DestroyCell.REASON_FOR_INT.get(reason, CloseReason.UNKNOWN)
-      self.reason_int = reason
-    elif reason in CloseReason:
-      self.reason = reason
-      self.reason_int = DestroyCell.INT_FOR_REASON.get(reason, -1)
-    else:
-      raise ValueError('Invalid closure reason: %s' % reason)
+    self.reason, self.reason_int = CloseReason.get(reason)
 
   @classmethod
   def pack(cls, link_version, circ_id, reason = CloseReason.NONE):
@@ -340,12 +294,7 @@ class DestroyCell(CircuitCell):
     :returns: **bytes** to close the circuit
     """
 
-    reason = DestroyCell.INT_FOR_REASON.get(reason, reason)
-
-    if not isinstance(reason, int):
-      raise ValueError('Invalid closure reason: %s' % reason)
-
-    return cls._pack(link_version, Size.CHAR.pack(reason), circ_id)
+    return cls._pack(link_version, Size.CHAR.pack(CloseReason.get(reason)[1]), circ_id)
 
   @classmethod
   def _unpack(cls, content, circ_id, link_version):



_______________________________________________
tor-commits mailing list
tor-commits@xxxxxxxxxxxxxxxxxxxx
https://lists.torproject.org/cgi-bin/mailman/listinfo/tor-commits