Refactored database, added different ways of retrieving pokemon & made unit test class

UnitTest.py

@@ -0,0 +1,73 @@
+# This module is for testing the different components of this project.
+
+from database import Database
+from sys import argv
+
+
+def print_items(items):
+    # Print each item in a collection.
+    for item in items:
+        print(item)
+
+
+def test_database_single_arg(arg):
+    # Test the database where there is a single command line parameter.
+    # The parameter is the name of the method to test.
+    arg = arg[1].lower()
+    db = Database()
+    if arg == "__str__":
+        print(db)
+    elif arg == "get_all":
+        print_items(db.get_all())
+    elif arg == "get_regions":
+        print_items(db.get_regions())
+    elif arg == "get_kanto":
+        print_items(db.get_kanto())
+    elif arg == "get_johto":
+        print_items(db.get_johto())
+    elif arg == "get_hoenn":
+        print_items(db.get_hoenn())
+    elif arg == "get_sinnoh":
+        print_items(db.get_sinnoh())
+    elif arg == "get_extra":
+        print_items(db.get_extra())
+    elif arg == "get_random":
+        print(db.get_random())
+    else:
+        print("No such public method '" + arg + "' with zero parameters exists in the Database class.")
+
+
+def test_database_double_arg(arg):
+    # Test the database where there are two command line parameters.
+    # The first parameter is the name of the method to test.
+    # The second parameter is the input parameter for the method that is being test.
+    arg1 = arg[1].lower()
+    arg2 = arg[2].lower()
+    db = Database()
+    if arg1 == "pokemon_exists":
+        print(db.pokemon_exists(arg2))
+    elif arg1 == "pokemon_id_exists":
+        print(db.pokemon_id_exists(arg2))
+    elif arg1 == "pokemon_name_exists":
+        print(db.pokemon_name_exists(arg2))
+    elif arg1 == "get_pokemon":
+        print(db.get_pokemon(arg2))
+    elif arg1 == "get_pokemon_by_name":
+        print(db.get_pokemon_by_name(arg2))
+    elif arg1 == "get_pokemon_by_id":
+        print(db.get_pokemon_by_id(arg2))
+    elif arg1 == "names_with_prefix":
+        print_items(db.names_with_prefix(arg2))
+    elif arg1 == "names_with_infix":
+        print_items(db.names_with_infix(arg2))
+
+
+if __name__ == "__main__":
+    if len(argv) == 1:
+        print("No command line parameters provided.")
+    elif len(argv) == 2:
+        test_database_single_arg(argv)
+    elif len(argv) == 3:
+        test_database_double_arg(argv)
+    else:
+        print("This module only takes one command line parameter.")

database.py

@@ -1,6 +1,7 @@
 # The Database object is a container for all the supported Pokemon.
 
-import os, random
+import os
+import random
 
 
 class Pokemon:
@@ -50,7 +51,7 @@ class Database:
         self.__load_extra()
 
     def __str__(self):
-        string = "POKEMON:\n"
+        string = ""
         for element in self.__pokemon_list:
             string += str(element) + "\n"
         return string[:-1]  # Remove the final new line ("\n").
@@ -62,6 +63,10 @@ class Database:
             result.append(pokemon)
         return result
 
+    def get_regions(self):
+        # Get all the supported regions.
+        return self.__regions
+
     def get_kanto(self):
         # Get all the Pokemon from the Kanto region.
         return self.__get_region("kanto")
@@ -82,10 +87,6 @@ class Database:
         # Get all the Extra Pokemon images available.
         return self.__get_region(None)
 
-    def get_regions(self):
-        # Get all the supported regions.
-        return self.__regions
-
     def __get_region(self, region):
         # Helper method for getting all the Pokemon of a specified region.
         result = []
@@ -94,14 +95,19 @@ class Database:
                 result.append(pokemon)
         return result
 
+    def get_random(self):
+        # Select a random Pokemon from the database.
+        random_int = random.randint(0, len(self.__pokemon_list))
+        return self.__pokemon_list[random_int]
+
     def pokemon_exists(self, pokemon):
         # Check for a Pokemon by ID or name.
         if type(pokemon) is int or str(pokemon).isdigit():
-            return self.id_exists(pokemon)
+            return self.pokemon_id_exists(pokemon)
         else:
-            return self.name_exists(pokemon)
+            return self.pokemon_name_exists(pokemon)
 
-    def id_exists(self, identifier):
+    def pokemon_id_exists(self, identifier):
         # Check for Pokemon by ID.
         identifier = int(identifier)
         if identifier < 1 or identifier > self.__MAX_ID:
@@ -109,22 +115,52 @@ class Database:
         else:
             return True
 
-    def name_exists(self, name):
+    def pokemon_name_exists(self, name):
         # Check for Pokemon by Name.
         return name.lower() in self.__pokemon_dictionary
 
-    def names_starting_with(self, prefix):
-        # Return Pokemon who's names contain the specified prefix.
+    def get_pokemon(self, pokemon):
+        # Get a Pokemon by name or ID.
+        if type(pokemon) is not int and type(pokemon) is not str:
+            raise Exception("The parameter Pokemon must be of type integer or string.")
+        if not self.pokemon_exists(pokemon):
+            raise Exception("No such Pokemon in the database.")
+        if type(pokemon) is int or str(pokemon).isdigit():
+            return self.get_pokemon_by_id(pokemon)
+        else:
+            return self.get_pokemon_by_name(pokemon)
+
+    def get_pokemon_by_name(self, name):
+        # Get a Pokemon by its name.
+        if type(name) is not str:
+            raise TypeError("The type of name must be a string.")
+        if not self.pokemon_name_exists(name):
+            raise Exception("No such Pokemon in the database.")
+        return self.__pokemon_dictionary[name]
+
+    def get_pokemon_by_id(self, identifier):
+        # Get a Pokemon by its ID.
+        if type(identifier) is not int and not str(identifier).isdigit():
+            raise TypeError("The Pokemon ID must be a number.")
+        if not self.pokemon_id_exists(identifier):
+            raise Exception("The Pokemon ID must be between 1 and " + str(self.__MAX_ID) + " inclusive.")
+        return self.__pokemon_list[int(identifier) - 1]  # Subtract 1 to convert to 0 base indexing.
+
+    def names_with_prefix(self, prefix):
+        # Return Pokemon who's names begin with the specified prefix.
         result = []
         for pokemon in self.__pokemon_list:
             if str(pokemon.get_name()).startswith(prefix):
                 result.append(pokemon)
         return result
 
-    def get_random(self):
-        # Select a random Pokemon from the database.
-        random_int = random.randint(1, len(self.__pokemon_list))
-        return self.__pokemon_list[random_int]
+    def names_with_infix(self, infix):
+        # Return Pokemon who's names contains the specified infix.
+        result = []
+        for pokemon in self.__pokemon_list:
+            if infix in str(pokemon.get_name()):
+                result.append(pokemon)
+        return result
 
     def __load_data(self):
         # Load all the Pokemon data. This does not include the 'Extra' Pokemon.

main.py

@@ -4,6 +4,7 @@
 
 from sys import argv
 from database import Pokemon, Database
+import scripter
 
 
 def print_list(list_of_items):
@@ -37,7 +38,7 @@ def print_columns(items):
 
 def prefix_search(db, arg):
     # Find all Pokemon in database, db, with the prefix, arg.
-    result = db.names_starting_with(arg)
+    result = db.names_with_prefix(arg)
     if len(result) == 0:
         print("No Pokemon found with prefix '" + arg + "'.")
     else:
@@ -102,6 +103,8 @@ def single_argument_handler(arg):
         print_columns(db.get_all())
     elif arg == "rand" or arg == "random":
         print(db.get_random())
+    else:
+        scripter.change_terminal()
 
 if __name__ == "__main__":
     # Entrance to the program.