diff --git a/school_center.py b/school_center.py index 40b2f66..89ed29c 100644 --- a/school_center.py +++ b/school_center.py @@ -1,7 +1,6 @@ from utils.custom_logger import configure_logging from typing import Dict, List -import sys -import os +from os import sys, path, makedirs import argparse import logging import random @@ -14,19 +13,13 @@ MIN_STUDENT_IN_CENTER = 10 # Min. no of students from a school to be assigned to a center in normal circumstances STRETCH_CAPACITY_FACTOR = 0.02 # How much can center capacity be streched if need arises PREF_CUTOFF = -4 # Do not allocate students with pref score less than cutoff +DEFAULT_OUTPUT_DIR = 'results' # Default directory to create output files if --output not provided +DEFAULT_OUTPUT_FILENAME = 'school-center.tsv' configure_logging() logger = logging.getLogger(__name__) -def create_dir(dirPath: str): - """ - Create the given directory if it doesn't exists - - Creates all the directories needed to resolve to the provided directory path - """ - if not os.path.exists(dirPath): - os.makedirs(dirPath) - def haversine_distance(lat1, lon1, lat2, lon2): """ @@ -214,7 +207,7 @@ def is_allocated(scode1: str, scode2: str) -> bool: parser.add_argument('prefs_tsv', default='prefs.tsv', help="Tab separated (TSV) file containing preference scores") parser.add_argument( - '-o', '--output', default='school-center.tsv', help='Output file') + '-o', '--output', default = DEFAULT_OUTPUT_FILENAME, help='Output file') parser.add_argument('-s', '--seed', action='store', metavar='SEEDVALUE', default=None, type=float, help='Initialization seed for Random Number Generator') @@ -231,10 +224,28 @@ def is_allocated(scode1: str, scode2: str) -> bool: remaining = 0 # stores count of non allocated students allocations = {} # to track mutual allocations -OUTPUT_DIR = 'results/' -create_dir(OUTPUT_DIR) # Create the output directory if not exists -with open('{}school-center-distance.tsv'.format(OUTPUT_DIR), 'w', encoding='utf-8') as intermediate_file, \ - open(OUTPUT_DIR + args.output, 'w', encoding='utf-8') as a_file: + +def get_output_dir(): + dirname = path.dirname(args.output) + if(dirname): + return dirname + else: + return DEFAULT_OUTPUT_DIR + +def get_output_filename(): + basename = path.basename(args.output) + if(basename): + return basename + else: + return DEFAULT_OUTOUT_FILENAME + + +output_dirname = get_output_dir() +output_filename = get_output_filename() +makedirs(output_dirname, exist_ok=True) # Create the output directory if not exists + +with open(path.join(output_dirname, "school-center-distance.tsv"), 'w', encoding='utf-8') as intermediate_file, \ +open(path.join(output_dirname, output_filename), 'w', encoding='utf-8') as a_file: writer = csv.writer(intermediate_file, delimiter="\t") writer.writerow(["scode", "s_count",