Alter DeviceInfoOrigin to Enum

app/collect/database.py

@@ -25,7 +25,7 @@ def upsert(session, model, rows, update_cols):
 
 
 def update_device_infos(session, address_origin, path=None):
-    if address_origin == DeviceInfoOrigin.flarmnet:
+    if address_origin == DeviceInfoOrigin.FLARMNET:
         device_infos = get_flarmnet(fln_file=path)
     else:
         device_infos = get_ddb(csv_file=path)
@@ -49,7 +49,7 @@ def import_ddb(session, logger=None):
         logger = app.logger
 
     logger.info("Import registered devices fom the DDB...")
-    counter = update_device_infos(session, DeviceInfoOrigin.ogn_ddb)
+    counter = update_device_infos(session, DeviceInfoOrigin.OGN_DDB)
 
     finish_message = "DeviceInfo: {} inserted.".format(counter)
     logger.info(finish_message)

app/commands/database.py

@@ -94,7 +94,7 @@ def import_ddb():
     """Import registered devices from the DDB."""
 
     print("Import registered devices fom the DDB...")
-    counter = update_device_infos(db.session, DeviceInfoOrigin.ogn_ddb)
+    counter = update_device_infos(db.session, DeviceInfoOrigin.OGN_DDB)
     print("Imported %i devices." % counter)
 
 
@@ -104,7 +104,7 @@ def import_file(path="tests/custom_ddb.txt"):
     """Import registered devices from a local file."""
 
     print("Import registered devices from '{}'...".format(path))
-    counter = update_device_infos(db.session, DeviceInfoOrigin.user_defined, path=path)
+    counter = update_device_infos(db.session, DeviceInfoOrigin.USER_DEFINED, path=path)
     print("Imported %i devices." % counter)
 
 
@@ -114,7 +114,7 @@ def import_flarmnet(path=None):
     """Import registered devices from a local file."""
 
     print("Import registered devices from '{}'...".format("internet" if path is None else path))
-    counter = update_device_infos(db.session, DeviceInfoOrigin.flarmnet, path=path)
+    counter = update_device_infos(db.session, DeviceInfoOrigin.FLARMNET, path=path)
     print("Imported %i devices." % counter)
 
 

app/model/device_info.py

@@ -1,4 +1,5 @@
 from app import db
+from .device_info_origin import DeviceInfoOrigin
 
 
 class DeviceInfo(db.Model):
@@ -15,7 +16,7 @@ class DeviceInfo(db.Model):
     identified = db.Column(db.Boolean)
     aircraft_type = db.Column(db.SmallInteger)
 
-    address_origin = db.Column(db.SmallInteger)
+    address_origin = db.Column(db.Enum(DeviceInfoOrigin), nullable=False, default=DeviceInfoOrigin.UNKNOWN)
 
     def __repr__(self):
         return "<DeviceInfo: %s,%s,%s,%s,%s,%s,%s,%s,%s>" % (

app/model/device_info_origin.py

@@ -1,22 +1,7 @@
-class DeviceInfoOrigin:
-    unknown = 0
-    ogn_ddb = 1
-    flarmnet = 2
-    user_defined = 3
+import enum
 
-    def __init__(self, origin=0):
-        if origin in [0, 1, 2, 3]:
-            self.origin = origin
-        else:
-            raise ValueError("no address origin with id {} known".format(origin))
-
-    def name(self):
-        if self.origin == self.unknown:
-            return "unknown"
-        elif self.origin == self.ogn_ddb:
-            return "OGN-DDB"
-        elif self.origin == self.flarmnet:
-            return "FlarmNet"
-        elif self.origin == self.user_defined:
-            return "user-defined"
-        return ""
+class DeviceInfoOrigin(enum.Enum):
+    UNKNOWN = 0
+    OGN_DDB = 1
+    FLARMNET = 2
+    USER_DEFINED = 3

app/utils.py

@@ -31,7 +31,7 @@ def date_to_timestamps(date):
     return (start, end)
 
 
-def get_ddb(csv_file=None, address_origin=DeviceInfoOrigin.unknown):
+def get_ddb(csv_file=None, address_origin=DeviceInfoOrigin.UNKNOWN):
     if csv_file is None:
         r = requests.get(DDB_URL)
         rows = "\n".join(i for i in r.text.splitlines() if i[0] != "#")
@@ -59,7 +59,7 @@ def get_ddb(csv_file=None, address_origin=DeviceInfoOrigin.unknown):
     return device_infos
 
 
-def get_flarmnet(fln_file=None, address_origin=DeviceInfoOrigin.flarmnet):
+def get_flarmnet(fln_file=None, address_origin=DeviceInfoOrigin.FLARMNET):
     if fln_file is None:
         r = requests.get(FLARMNET_URL)
         rows = [bytes.fromhex(line).decode("latin1") for line in r.text.split("\n") if len(line) == 172]

migrations/versions/be9a6dad551e_use_enums_for_deviceinfoorigin.py

@@ -0,0 +1,41 @@
+"""Use Enums for DeviceInfoOrigin
+
+Revision ID: be9a6dad551e
+Revises: 885123e6a2d6
+Create Date: 2019-09-15 14:38:25.838089
+
+"""
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql    
+
+
+# revision identifiers, used by Alembic.
+revision = 'be9a6dad551e'
+down_revision = '885123e6a2d6'
+branch_labels = None
+depends_on = None
+
+deviceinfoorigin = postgresql.ENUM('UNKNOWN', 'OGN_DDB', 'FLARMNET', 'USER_DEFINED', name='deviceinfoorigin')
+
+def upgrade():
+    deviceinfoorigin.create(op.get_bind())
+    
+    op.add_column('device_infos', sa.Column('address_origin_enum', sa.Enum('UNKNOWN', 'OGN_DDB', 'FLARMNET', 'USER_DEFINED', name='deviceinfoorigin'), nullable=False, server_default='UNKNOWN'))
+    op.execute("UPDATE device_infos SET address_origin_enum = 'UNKNOWN' WHERE address_origin = 0")
+    op.execute("UPDATE device_infos SET address_origin_enum = 'OGN_DDB' WHERE address_origin = 1")
+    op.execute("UPDATE device_infos SET address_origin_enum = 'FLARMNET' WHERE address_origin = 2")
+    op.execute("UPDATE device_infos SET address_origin_enum = 'USER_DEFINED' WHERE address_origin = 3")
+    op.drop_column('device_infos', 'address_origin')
+    op.alter_column('device_infos', 'address_origin_enum', new_column_name='address_origin')
+
+def downgrade():
+    op.add_column('device_infos', sa.Column('address_origin_int', sa.SmallInteger))
+    op.execute("UPDATE device_infos SET address_origin_int = 0 WHERE address_origin = 'UNKNOWN'")
+    op.execute("UPDATE device_infos SET address_origin_int = 1 WHERE address_origin = 'OGN_DDB'")
+    op.execute("UPDATE device_infos SET address_origin_int = 2 WHERE address_origin = 'FLARMNET'")
+    op.execute("UPDATE device_infos SET address_origin_int = 3 WHERE address_origin = 'USER_DEFINED'")
+    op.drop_column('device_infos', 'address_origin')
+    op.alter_column('device_infos', 'address_origin_int', new_column_name='address_origin')
+
+    deviceinfoorigin.drop(op.get_bind())