Storage get count

AUTHORS

@@ -3,4 +3,5 @@
 
 Jennifer Huang <133@holbertonschool.com>
 Alexa Orrico <210@holbertonschool.com>
-Joann Vuong <130@holbertonschool.com> 
+Joann Vuong <130@holbertonschool.com>
+Mark Manani <mikonani413@gmail.com>

models/engine/db_storage.py

@@ -15,6 +15,7 @@
 import sqlalchemy
 from sqlalchemy import create_engine
 from sqlalchemy.orm import scoped_session, sessionmaker
+from sqlalchemy.exc import SQLAlchemyError
 
 classes = {"Amenity": Amenity, "City": City,
            "Place": Place, "Review": Review, "State": State, "User": User}
@@ -74,3 +75,20 @@ def reload(self):
     def close(self):
         """call remove() method on the private session attribute"""
         self.__session.remove()
+
+    def get(self, cls, id):
+        """retrieves one object"""
+        try:
+            return self.__session.query(cls).get(id)
+        except SQLAlchemyError:
+            return None
+
+    def count(self, cls=None):
+        """Count the number of objects in storage for the given class."""
+        try:
+            if cls:
+                return self.__session.query(cls).count()
+            else:
+                return sum([self.__session.query(c).count() for c in self.__classes.values()])
+        except SQLAlchemyError:
+            return 0

models/engine/file_storage.py

@@ -68,3 +68,14 @@ def delete(self, obj=None):
     def close(self):
         """call reload() method for deserializing the JSON file to objects"""
         self.reload()
+
+    def get(self, cls, id):
+        """Retrieve one object based on the class and its ID."""
+        objects = self.all(cls)
+        return objects.get(id, None)
+
+    def count(self, cls=None):
+        """Count the number of objects in storage for the given class."""
+        if cls:
+            return len([obj for obj in self.__objects.values() if isinstance(obj, cls)])
+        return len(self.__objects)

tests/test_models/test_amenity.py

@@ -8,7 +8,7 @@
 import models
 from models import amenity
 from models.base_model import BaseModel
-import pep8
+import pep8 as pycodestyle
 import unittest
 Amenity = amenity.Amenity
 

tests/test_models/test_city.py

@@ -8,7 +8,7 @@
 import models
 from models import city
 from models.base_model import BaseModel
-import pep8
+import pep8 as pycodestyle
 import unittest
 City = city.City