From e4cfd964d593088511506014ca77c71f20126276 Mon Sep 17 00:00:00 2001
From: Lonami Exo <totufals@hotmail.com>
Date: Mon, 3 Dec 2018 15:56:38 +0100
Subject: [PATCH] Use an enum for methods' usability

---
 telethon_generator/parsers/methods.py | 22 ++++++++++++++++++----
 1 file changed, 18 insertions(+), 4 deletions(-)

diff --git a/telethon_generator/parsers/methods.py b/telethon_generator/parsers/methods.py
index 6d2572a0..fef06953 100644
--- a/telethon_generator/parsers/methods.py
+++ b/telethon_generator/parsers/methods.py
@@ -1,11 +1,28 @@
 import csv
+import enum
+
+
+class Usability(enum.Enum):
+    UNKNOWN = 0
+    USER = 1
+    BOT = 2
+    BOTH = 4
 
 
 class MethodInfo:
     def __init__(self, name, usability, errors):
         self.name = name
-        self.usability = usability
         self.errors = errors
+        try:
+            self.usability = {
+                'unknown': Usability.UNKNOWN,
+                'user': Usability.USER,
+                'bot': Usability.BOT,
+                'both': Usability.BOTH,
+            }[usability.lower()]
+        except KeyError:
+            raise ValueError('Usability must be either user, bot, both or '
+                             'unknown, not {}'.format(usability)) from None
 
 
 def parse_methods(csv_file, errors_dict):
@@ -17,9 +34,6 @@ def parse_methods(csv_file, errors_dict):
         f = csv.reader(f)
         next(f, None)  # header
         for line, (method, usability, errors) in enumerate(f, start=2):
-            if usability not in ('user', 'bot', 'both', 'unknown'):
-                raise ValueError('Usability must be either user, bot, '
-                                 'both or unknown, not {}'.format(usability))
             try:
                 errors = [errors_dict[x] for x in errors.split()]
             except KeyError: