lo-ranking/base.py

from __future__ import print_function
import csv
import os

RAW_DIR = os.path.join(os.getcwd(), 'raw', 'july')
TARGET_DIR = os.path.join(os.getcwd(), 'build')
MODEL_LIST = [
    'PMSf', 'PMSD', 'PMHf', 'PMHD', 'PFSf',
    'PFSD', 'PFHf', 'PFHD', 'IMSf', 'IMSD',
    'IMHf', 'IMHD', 'IFSf', 'IFSD', 'IFHf',
    'IFHD', 'CMSf', 'CMSD', 'CMHf', 'CMHD',
    'CFSf', 'CFSD', 'CFHf', 'CFHD'
]


class Lo:
    src_path = None
    los = {}
    model = {
        'A': {}, 'B': {}, 'C': {}, 'D': {},
        'E': {}, 'F': {}, 'G': {}, 'H': {},
    }

    def __init__(self, filename=None):
        base_path = os.getcwd()
        f_name = '03_Lo_data.csv' if not filename else filename
        self.src_path = os.path.join(base_path, 'raw', f_name)
        self.los = {}
        with open(self.src_path, 'r') as f:
            reader = csv.reader(f)
            reader.next()
            for r in reader:
                lo_id = r[0]
                t = {
                    'A': '%s%s%s%s' % (r[2], r[3], r[6], r[7]),
                    'B': '%s%s%s' % (r[10], r[12], r[13]),
                    'C': '%s%s%s%s' % (r[15], r[17], r[18], r[19]),
                    'D': '%s%s%s' % (r[21], r[22], r[24]),
                    'E': '%s%s%s%s%s' % (r[27], r[29], r[30], r[31], r[33]),
                    'F': '%s%s%s' % (r[13], r[15], r[19]),
                    'G': '%s%s%s' % (r[7], r[13], r[21]),
                    'H': '%s%s%s' % (r[10], r[22], r[27]),
                }
                for k, v in t.items():
                    if v not in self.model[k]:
                        self.model[k][v] = []
                    self.model[k][v].append(lo_id)
                self.los[lo_id] = t.copy()
                self.los[lo_id]['weight'] = r[35]
                self.los[lo_id]['model'] = r[1]

    def get(self, lo_id):
        try:
            return self.los[lo_id]
        except (IndexError, KeyError):
            return {}

    def get_lo_list(self):
        return self.los.keys()

    def get_model(self, m_name):
        try:
            return self.model[m_name]
        except IndexError:
            return []

    def get_mm(self, m_name, value):
        return self.get_model_match(m_name, value)

    def get_model_match(self, m_name, value):
        try:
            return self.model[m_name][value]
        except IndexError:
            return []

    def get_mm_tuple(self, *args):
        if len(args) > 1 or not isinstance(args, tuple):
            return []
        r = []
        for i in args[0]:
            if len(i) != 2 or not i[1]:
                continue
            m = self.get_mm(i[0], i[1])
            if m:
                r.append(m)
        return set.intersection(*[set(k) for k in r])


def create_csv(filename, data, **kwargs):
    if not isinstance(data, list):
        print('Nope, data are not list')
    directory = kwargs.get('directory', '')
    path = os.path.join(TARGET_DIR, directory)
    try:
        os.mkdir(path)
    except OSError:
        pass
    fpath = os.path.join(path, filename)
    with open(fpath, 'wb') as f:
        owriter = csv.writer(
            f, delimiter=',', quotechar='"', quoting=csv.QUOTE_ALL)
        for r in data:
            owriter.writerow(r)