refactor functions

util/course-permutation-tool/generate_permutations.py

@@ -13,17 +13,17 @@ import pytz
 import sys
 
 # course defaults
-course_number = None
-course_organization = "RITX"
-course_run = "3T2017"
-course_user = "edx@example.com"
-course_partner = "edx"
+COURSE_NUMBER = None
+COURSE_ORGANIZATION = "RITX"
+COURSE_RUN = "3T2017"
+COURSE_USER = "edx@example.com"
+COURSE_PARTNER = "edx"
 
 # date configurations
-future_date = "future"
-past_date = "past"
+FUTURE_DATE = "future"
+PAST_DATE = "past"
 
-test_courses_file = "test_courses.json"
+TEST_COURSES_FILE = "test_courses.json"
 
 
 def arg_parse():
@@ -35,7 +35,7 @@ def arg_parse():
     return parser.parse_args()
 
 
-def process_field_arguments(parser_args, seed_data_file, field_args):
+def process_field_arguments(parser_args):
     try:
         file = open(parser_args.seed_data_file)
         file_data = json.load(file)
@@ -43,31 +43,32 @@ def process_field_arguments(parser_args, seed_data_file, field_args):
         permutation_data = file_data["permutation_data"]
         default_data_keys = permutation_data.keys()
     except IOError:
-        print sys.exc_value
+        print("Please fix your input file - {}".format(sys.exc_value))
         sys.exit()
     except KeyError:
         print("Could not find key {}".format(sys.exc_value))
         sys.exit()
 
     field_args = {}
+    parsed_field_args = parser_args.fields
 
     # if no field arguments are given, just return default data
-    if not parser_args.fields:
+    if not parsed_field_args:
         field_args = file_data["default_data"]
     else:
-        num_field_args = len(parser_args.fields[0])
+        num_field_args = len(parsed_field_args[0])
         if (num_field_args > 3):
             raise argparse.ArgumentTypeError("Only a max of 3 fields allowed")
         # add each command line field to fields dict
-        for permutation_choices in parser_args.fields:
-            for i in range(0, num_field_args):
-                try:
-                    field_args[permutation_choices[i]] = permutation_data[permutation_choices[i]]
-                except KeyError:
-                    print "{} is not a field option".format(sys.exc_value)
+        for field in parsed_field_args[0]:
+            try:
+                field_args[field] = permutation_data[field]
+            except KeyError:
+                print "{} is not a field option".format(sys.exc_value)
+                sys.exit()
 
-        # the difference btwn all possible fields and the chosen permutation ones
-        default_fields = list(set(default_data_keys) - set(parser_args.fields[0]))
+        # calculate the difference between all possible fields and the chosen permutation ones
+        default_fields = list(set(default_data_keys) - set(parsed_field_args[0]))
 
         # add non permutation fields to dict
         for field in default_fields:
@@ -82,8 +83,19 @@ def generate_permutations(field_args, index, results, courses_dict, field_values
     permutations_values = field_args[permutation_option]
 
     for permutation_value in permutations_values:
-        # add other required default fields to dict
-        enrollment_dict = set_course_defaults(courses_dict)
+        # add other required course and enrollment default fields to dict
+        # will be generated automatically by course creation script
+        courses_dict["number"] = COURSE_NUMBER
+        courses_dict["organization"] = COURSE_ORGANIZATION
+        courses_dict["run"] = COURSE_RUN
+        courses_dict["user"] = COURSE_USER
+        courses_dict["partner"] = COURSE_PARTNER
+        # configure enrollment seat settings
+        enrollment_dict = {}
+        enrollment_dict["credit"] = False
+        enrollment_dict["credit_provider"] = "test-credit-provider"
+
+
         # add permutation fields to dict
         field_values_dict[permutation_option] = permutation_value
         # generate start and end dates
@@ -106,29 +118,20 @@ def generate_permutations(field_args, index, results, courses_dict, field_values
     wrapper_courses_dict = {}  # needed to match course input file creation
     wrapper_courses_dict["courses"] = results
 
-    with open(test_courses_file, "w") as outfile:
-        json.dump(wrapper_courses_dict, outfile)
+    create_courses_json_file(wrapper_courses_dict)
 
 
-def set_course_defaults(courses_dict):
-    courses_dict["number"] = course_number  # will be generated automatically by course creation script
-    courses_dict["organization"] = course_organization
-    courses_dict["run"] = course_run
-    courses_dict["user"] = course_user
-    courses_dict["partner"] = course_partner
-    # configure enrollment seat settings
-    enrollment_dict = {}
-    enrollment_dict["credit"] = False
-    enrollment_dict["credit_provider"] = "test-credit-provider"
-    return enrollment_dict
+def create_courses_json_file(wrapper_courses_dict):
+    with open(TEST_COURSES_FILE, "w") as outfile:
+        json.dump(wrapper_courses_dict, outfile)
 
 
 def generate_date_translation(field_values_dict, permutation_option, permutation_value):
     now = datetime.datetime.now(pytz.UTC)
-    if permutation_value == future_date:
+    if permutation_value == FUTURE_DATE:
         future = str(now + datetime.timedelta(days=365))
         field_values_dict[permutation_option] = future
-    if permutation_value == past_date:
+    if permutation_value == PAST_DATE:
         past = str(now - datetime.timedelta(days=60))
         field_values_dict[permutation_option] = past
     if permutation_value == None:
@@ -141,5 +144,5 @@ def start_field_recursion(process_field_args):
 
 if __name__ == "__main__":
     args = arg_parse()
-    process_field_args = process_field_arguments(args, args.seed_data_file, args.fields)
+    process_field_args = process_field_arguments(args)
     start_field_recursion(process_field_args)