refactored object getting methods in the database class

pycs/database/Database.py

@@ -155,16 +155,47 @@ class Database:
     def __exit__(self, exc_type, exc_val, exc_tb):
         self.con.__exit__(exc_type, exc_val, exc_tb)
 
+    def get_object_by_id(self, table_name: str, identifier: int, cls):
+        """
+        create an object from cls and a row fetched from table_name and
+        identified by the identifier
+
+        :param table_name: table name
+        :param identifier: unique identifier
+        :param cls: class that is used to create the object
+        :return: object of type cls
+        """
+        with closing(self.con.cursor()) as cursor:
+            cursor.execute(f'SELECT * FROM {table_name} WHERE id = ?', [identifier])
+            row = cursor.fetchone()
+
+            if row is not None:
+                return cls(self, row)
+
+            return None
+
+    def get_objects(self, table_name: str, cls):
+        """
+        get a list of all available objects in the table
+
+        :param table_name: table name
+        :param cls: class that is used to create the objects
+        :return: list of object of type cls
+        """
+
+        with closing(self.con.cursor()) as cursor:
+            cursor.execute(f'SELECT * FROM {table_name}')
+            for row in cursor:
+                yield cls(self, row)
+
+
     def models(self) -> Iterator[Model]:
         """
         get a list of all available models
 
         :return: iterator of models
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM models')
-            for row in cursor:
-                yield Model(self, row)
+        return self.get_objects("models", Model)
 
     def model(self, identifier: int) -> Optional[Model]:
         """
@@ -173,14 +204,7 @@ class Database:
         :param identifier: unique identifier
         :return: model
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM models WHERE id = ?', [identifier])
-            row = cursor.fetchone()
-
-            if row is not None:
-                return Model(self, row)
-
-            return None
+        return self.get_object_by_id("models", identifier, Model)
 
     def label_providers(self) -> Iterator[LabelProvider]:
         """
@@ -188,10 +212,7 @@ class Database:
 
         :return: iterator over label providers
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM label_providers')
-            for row in cursor:
-                yield LabelProvider(self, row)
+        return self.get_objects("label_providers", LabelProvider)
 
     def label_provider(self, identifier: int) -> Optional[LabelProvider]:
         """
@@ -200,14 +221,7 @@ class Database:
         :param identifier: unique identifier
         :return: label provider
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM label_providers WHERE id = ?', [identifier])
-            row = cursor.fetchone()
-
-            if row is not None:
-                return LabelProvider(self, row)
-
-            return None
+        return self.get_object_by_id("label_providers", identifier, LabelProvider)
 
     def projects(self) -> Iterator[Project]:
         """
@@ -215,10 +229,7 @@ class Database:
 
         :return: iterator over projects
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM projects')
-            for row in cursor:
-                yield Project(self, row)
+        return self.get_objects("projects", Project)
 
     def project(self, identifier: int) -> Optional[Project]:
         """
@@ -227,14 +238,7 @@ class Database:
         :param identifier: unique identifier
         :return: project
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM projects WHERE id = ?', [identifier])
-            row = cursor.fetchone()
-
-            if row is not None:
-                return Project(self, row)
-
-            return None
+        return self.get_object_by_id("projects", identifier, Project)
 
     def create_project(self,
                        name: str,
@@ -280,14 +284,7 @@ class Database:
         :param identifier: unique identifier
         :return: collection
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM collections WHERE id = ?', [identifier])
-            row = cursor.fetchone()
-
-            if row is not None:
-                return Collection(self, row)
-
-            return None
+        return self.get_object_by_id("collections", identifier, Collection)
 
     def file(self, identifier) -> Optional[File]:
         """
@@ -296,14 +293,7 @@ class Database:
         :param identifier: unique identifier
         :return: file
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM files WHERE id = ?', [identifier])
-            row = cursor.fetchone()
-
-            if row is not None:
-                return File(self, row)
-
-            return None
+        return self.get_object_by_id("files", identifier, File)
 
     def result(self, identifier) -> Optional[Result]:
         """
@@ -312,11 +302,4 @@ class Database:
         :param identifier: unique identifier
         :return: result
         """
-        with closing(self.con.cursor()) as cursor:
-            cursor.execute('SELECT * FROM results WHERE id = ?', [identifier])
-            row = cursor.fetchone()
-
-            if row is not None:
-                return Result(self, row)
-
-            return None
+        return self.get_object_by_id("results", identifier, Result)