<commit_before>"""PyDoIt file for automating tasks."""

import glob
import itertools
import os
from os.path import dirname

from doit import get_var
#from doit.tools import Interactive
from doit.tools import create_folder
#from doit.tools import result_dep

import dwi.util

"""
Backends:
dbm: (default) It uses python dbm module.
json: Plain text using a json structure, it is slow but good for debugging.
sqlite3: (experimental) very slow implementation, support concurrent access.
"""

DOIT_CONFIG = {
    'backend': 'sqlite3',
    'default_tasks': [],
    'verbosity': 2,
    }

DWILIB = '~/src/dwilib/tools'
PMAP = DWILIB+'/pmap.py'
ANON = DWILIB+'/anonymize_dicom.py'
FIND_ROI = DWILIB+'/find_roi.py'
COMPARE_MASKS = DWILIB+'/compare_masks.py'
SELECT_VOXELS = DWILIB+'/select_voxels.py'
CALC_AUC = DWILIB+'/draw_roc.py'
CORRELATION = DWILIB+'/correlation.py'

MODELS = 'Si SiN Mono MonoN Kurt KurtN Stretched StretchedN '\
        'Biexp BiexpN'.split()

SAMPLELIST = get_var('samplelist', 'all') # Sample list (train, test, etc)
SAMPLELIST_FILE = 'samples_%s.txt' % SAMPLELIST
SAMPLES = dwi.util.read_sample_list(SAMPLELIST_FILE)
SUBWINDOWS = dwi.util.read_subwindows('subwindows.txt')

FIND_ROI_PARAMS = [
        [5], # ROI side min
        #[3, 5], # ROI side min
        [10, 15], # ROI side max
        #[10], # ROI side max
        range(500, 6000, 500), # Number of ROIs
]
#NROIS = [500, 1000, 1500, 2000, 2500, 3000, 3500, 4000, 4500]
#NROIS = FIND_ROI_PARAMS[-1]

# Common functions

def subwindow_to_str(subwindow):
    return ' '.join(map(str, subwindow))

def fit_cmd(model, subwindow, infiles, outfile):
    d = dict(prg=PMAP,
            m=model,
            sw=subwindow_to_str(subwindow),
            i=' '.join(infiles),
            o=outfile,
            )
    s = '{prg} -m {m} -s {sw} -d {i} -o {o}'.format(**d)
    return s

# Tasks

#def task_anonymize():
#    """Anonymize imaging data."""
#    files = glob.glob('dicoms/*/DICOMDIR') + glob.glob('dicoms/*/DICOM/*')
#    files.sort()
#    for f in files:
#        cmd = '{prg} -v -i {f}'.format(prg=ANON, f=f)
#        yield {
#                'name': f,
#                'actions': [cmd],
#                'file_dep': [f],
#                }

def task_fit():
    """Fit models to imaging data."""
    for case, scan in SUBWINDOWS.keys():
        subwindow = SUBWINDOWS[(case, scan)]
        d = dict(c=case, s=scan)
        infiles = sorted(glob.glob('dicoms/{c}_*_hB_{s}*/DICOM/*'.format(**d)))
        if len(infiles) == 0:
            continue
        for model in MODELS:
            d['m'] = model
            outfile = 'pmaps/pmap_{c}_{s}_{m}.txt'.format(**d)
            cmd = fit_cmd(model, subwindow, infiles, outfile)
            yield {
                    'name': '{c}_{s}_{m}'.format(**d),
                    'actions': [(create_folder, [dirname(outfile)]),
                            cmd],
                    'file_dep': infiles,
                    'targets': [outfile],
                    'clean': True,
                    }

def get_task_find_roi(case, scan, algparams):
    d = dict(prg=FIND_ROI, sl=SAMPLELIST_FILE, c=case, s=scan,
            algparams=' '.join(algparams), algparams_='_'.join(algparams))
    maskpath = 'masks_auto_{algparams_}/{c}_{s}_auto.mask'.format(**d)
    figpath = 'find_roi_images_{algparams_}/{c}_{s}.png'.format(**d)
    d.update(mp=maskpath, fp=figpath)
    file_deps = [SAMPLELIST_FILE]
    file_deps += glob.glob('masks_prostate/{c}_*_{s}_*/*'.format(**d))
    file_deps += glob.glob('masks_rois/{c}_*_{s}_*/*'.format(**d))
    cmd = '{prg} --samplelist {sl} --cases {c} --scans {s} '\
            '--algparams {algparams} --outmask {mp} --outfig {fp}'.format(**d)
    return {
            'name': '{c}_{s}_{algparams_}'.format(**d),
            'actions': [(create_folder, [dirname(maskpath)]),
                        (create_folder, [dirname(figpath)]),
                    cmd],
            'file_dep': file_deps,
            'targets': [maskpath, figpath],
            'clean': True,
            }

def task_find_roi():
    """Find a cancer ROI automatically."""
    for sample in SAMPLES:
        case = sample['case']
        for scan in sample['scans']:
            for algparams in itertools.product(*FIND_ROI_PARAMS):
                yield get_task_find_roi(case, scan, map(str, algparams))

## Deprecated.
#def task_compare_masks():
#    """Compare ROI masks."""
#    for case, scan in SUBWINDOWS.keys():
#        subwindow = SUBWINDOWS[(case, scan)]
#        d = dict(prg=COMPARE_MASKS,
#                c=case,
#                s=scan,
#                w=subwindow_to_str(subwindow),
#                )
#        #file1 = 'masks/{c}_{s}_ca.mask'.format(**d)
#        file1 = glob.glob('masks/{c}_{s}_1_*.mask'.format(**d))[0]
#        file2 = 'masks_auto/{c}_{s}_auto.mask'.format(**d)
#        outfile = 'roi_comparison/{c}_{s}.txt'.format(**d)
#        d['f1'] = file1
#        d['f2'] = file2
#        d['o'] = outfile
#        cmd = '{prg} -s {w} {f1} {f2} > {o}'.format(**d)
#        if not os.path.exists(file1):
#            print 'Missing mask file %s' % file1
#            continue
#        yield {
#                'name': '{c}_{s}'.format(**d),
#                'actions': [(create_folder, [dirname(outfile)]),
#                        cmd],
#                'file_dep': [file1, file2],
#                'targets': [outfile],
#                'clean': True,
#                }

def get_task_select_roi(case, scan, model, param, masktype, algparams,
        subwindow=None):
    d = dict(c=case, s=scan, m=model, p=param, mt=masktype,
            algparams_='_'.join(algparams), sw=subwindow)
    maskpath = 'masks_{mt}_{algparams_}/{c}_{s}_{mt}.mask'.format(**d)
    s = 'results_{m}_combinedDICOM/{c}_*_{s}/{c}_*_{s}_{p}'.format(**d)
    inpath = glob.glob(s)[0]
    outpath = 'rois_{mt}_{algparams_}/{c}_x_x_{s}_{m}_{p}_{mt}.txt'.format(**d)
    args = [SELECT_VOXELS]
    #args += ['-v']
    if subwindow:
        args += ['-s %s' % subwindow_to_str(subwindow)]
    args += ['-m %s' % maskpath]
    args += ['-i "%s"' % inpath]
    args += ['-o "%s"' % outpath]
    cmd = ' '.join(args)
    return {
            'name': '{c}_{s}_{algparams_}'.format(**d),
            'actions': [(create_folder, [dirname(outpath)]),
                    cmd],
            'file_dep': [maskpath],
            'targets': [outpath],
            'clean': True,
            }

def task_select_roi_cancer():
    """Select cancer ROIs from the pmap DICOMs."""
    for sample in SAMPLES:
        for scan in sample['scans']:
            case = sample['case']
            masktype = 'cancer'
            yield get_task_select_roi(case, scan, 'Mono', 'ADCm', masktype,
                    ['default'])

def task_select_roi_auto():
    """Select automatic ROIs from the pmap DICOMs."""
    for sample in SAMPLES:
        for scan in sample['scans']:
            for algparams in itertools.product(*FIND_ROI_PARAMS):
                case = sample['case']
                masktype = 'auto'
                yield get_task_select_roi(case, scan, 'Mono', 'ADCm', masktype,
                        map(str, algparams))

def task_evaluate_autoroi():
    """Evaluate auto-ROI prediction ability by ROC AUC and correlation with
    Gleason score."""
    outfile = 'autoroi_evaluation.txt'
    d = dict(prg_auc=CALC_AUC, prg_corr=CORRELATION, o=outfile)
    cmds = ['echo `date` > {o}'.format(**d)]
    for algparams in itertools.product(*FIND_ROI_PARAMS):
        d['algparams_'] = '_'.join(map(str, algparams))
        d['i'] = 'rois_auto_{algparams_}'.format(**d)
        s = 'echo -n {algparams_} >> {o}'
        cmds.append(s.format(**d))
        s = r'echo -n \\t`{prg_auc} -s patients.txt -l score -g 3+3 -m {i}/* -a --autoflip` >> {o}'
        cmds.append(s.format(**d))
        s = r'echo -n \\t`{prg_corr} -s patients.txt -l score -g 3+3 3+4 -a -m {i}/*` >> {o}'
        cmds.append(s.format(**d))
        s = r'echo \\t`{prg_corr} -s patients.txt -l score -a -m {i}/*` >> {o}'
        cmds.append(s.format(**d))
    return {
            'actions': cmds,
            'task_dep': ['select_roi_auto'],
            'targets': [outfile],
            'clean': True,
            }
<commit_msg>Added samplelist id to some paths<commit_after>"""PyDoIt file for automating tasks."""

import glob
import itertools
import os
from os.path import dirname

from doit import get_var
#from doit.tools import Interactive
from doit.tools import create_folder
#from doit.tools import result_dep

import dwi.util

"""
Backends:
dbm: (default) It uses python dbm module.
json: Plain text using a json structure, it is slow but good for debugging.
sqlite3: (experimental) very slow implementation, support concurrent access.
"""

DOIT_CONFIG = {
    'backend': 'sqlite3',
    'default_tasks': [],
    'verbosity': 2,
    }

DWILIB = '~/src/dwilib/tools'
PMAP = DWILIB+'/pmap.py'
ANON = DWILIB+'/anonymize_dicom.py'
FIND_ROI = DWILIB+'/find_roi.py'
COMPARE_MASKS = DWILIB+'/compare_masks.py'
SELECT_VOXELS = DWILIB+'/select_voxels.py'
CALC_AUC = DWILIB+'/draw_roc.py'
CORRELATION = DWILIB+'/correlation.py'

MODELS = 'Si SiN Mono MonoN Kurt KurtN Stretched StretchedN '\
        'Biexp BiexpN'.split()

SAMPLELIST = get_var('samplelist', 'all') # Sample list (train, test, etc)
SAMPLELIST_FILE = 'samples_%s.txt' % SAMPLELIST
SAMPLES = dwi.util.read_sample_list(SAMPLELIST_FILE)
SUBWINDOWS = dwi.util.read_subwindows('subwindows.txt')

FIND_ROI_PARAMS = [
        [5], # ROI side min
        #[3, 5], # ROI side min
        [10, 15], # ROI side max
        #[10], # ROI side max
        range(500, 6000, 500), # Number of ROIs
]
#NROIS = [500, 1000, 1500, 2000, 2500, 3000, 3500, 4000, 4500]
#NROIS = FIND_ROI_PARAMS[-1]

# Common functions

def subwindow_to_str(subwindow):
    return ' '.join(map(str, subwindow))

def fit_cmd(model, subwindow, infiles, outfile):
    d = dict(prg=PMAP,
            m=model,
            sw=subwindow_to_str(subwindow),
            i=' '.join(infiles),
            o=outfile,
            )
    s = '{prg} -m {m} -s {sw} -d {i} -o {o}'.format(**d)
    return s

# Tasks

#def task_anonymize():
#    """Anonymize imaging data."""
#    files = glob.glob('dicoms/*/DICOMDIR') + glob.glob('dicoms/*/DICOM/*')
#    files.sort()
#    for f in files:
#        cmd = '{prg} -v -i {f}'.format(prg=ANON, f=f)
#        yield {
#                'name': f,
#                'actions': [cmd],
#                'file_dep': [f],
#                }

def task_fit():
    """Fit models to imaging data."""
    for case, scan in SUBWINDOWS.keys():
        subwindow = SUBWINDOWS[(case, scan)]
        d = dict(c=case, s=scan)
        infiles = sorted(glob.glob('dicoms/{c}_*_hB_{s}*/DICOM/*'.format(**d)))
        if len(infiles) == 0:
            continue
        for model in MODELS:
            d['m'] = model
            outfile = 'pmaps/pmap_{c}_{s}_{m}.txt'.format(**d)
            cmd = fit_cmd(model, subwindow, infiles, outfile)
            yield {
                    'name': '{c}_{s}_{m}'.format(**d),
                    'actions': [(create_folder, [dirname(outfile)]),
                            cmd],
                    'file_dep': infiles,
                    'targets': [outfile],
                    'clean': True,
                    }

def get_task_find_roi(case, scan, algparams):
    d = dict(prg=FIND_ROI, sl=SAMPLELIST, slf=SAMPLELIST_FILE, c=case, s=scan,
            algparams=' '.join(algparams), algparams_='_'.join(algparams))
    maskpath = 'masks_auto_{sl}_{algparams_}/{c}_{s}_auto.mask'.format(**d)
    figpath = 'find_roi_images_{sl}_{algparams_}/{c}_{s}.png'.format(**d)
    d.update(mp=maskpath, fp=figpath)
    file_deps = [SAMPLELIST_FILE]
    file_deps += glob.glob('masks_prostate/{c}_*_{s}_*/*'.format(**d))
    file_deps += glob.glob('masks_rois/{c}_*_{s}_*/*'.format(**d))
    cmd = '{prg} --samplelist {slf} --cases {c} --scans {s} '\
            '--algparams {algparams} --outmask {mp} --outfig {fp}'.format(**d)
    return {
            'name': '{c}_{s}_{algparams_}'.format(**d),
            'actions': [(create_folder, [dirname(maskpath)]),
                        (create_folder, [dirname(figpath)]),
                    cmd],
            'file_dep': file_deps,
            'targets': [maskpath, figpath],
            'clean': True,
            }

def task_find_roi():
    """Find a cancer ROI automatically."""
    for sample in SAMPLES:
        case = sample['case']
        for scan in sample['scans']:
            for algparams in itertools.product(*FIND_ROI_PARAMS):
                yield get_task_find_roi(case, scan, map(str, algparams))

## Deprecated.
#def task_compare_masks():
#    """Compare ROI masks."""
#    for case, scan in SUBWINDOWS.keys():
#        subwindow = SUBWINDOWS[(case, scan)]
#        d = dict(prg=COMPARE_MASKS,
#                c=case,
#                s=scan,
#                w=subwindow_to_str(subwindow),
#                )
#        #file1 = 'masks/{c}_{s}_ca.mask'.format(**d)
#        file1 = glob.glob('masks/{c}_{s}_1_*.mask'.format(**d))[0]
#        file2 = 'masks_auto/{c}_{s}_auto.mask'.format(**d)
#        outfile = 'roi_comparison/{c}_{s}.txt'.format(**d)
#        d['f1'] = file1
#        d['f2'] = file2
#        d['o'] = outfile
#        cmd = '{prg} -s {w} {f1} {f2} > {o}'.format(**d)
#        if not os.path.exists(file1):
#            print 'Missing mask file %s' % file1
#            continue
#        yield {
#                'name': '{c}_{s}'.format(**d),
#                'actions': [(create_folder, [dirname(outfile)]),
#                        cmd],
#                'file_dep': [file1, file2],
#                'targets': [outfile],
#                'clean': True,
#                }

def get_task_select_roi(case, scan, model, param, masktype, algparams,
        subwindow=None):
    d = dict(sl=SAMPLELIST, c=case, s=scan, m=model, p=param, mt=masktype,
            algparams_='_'.join(algparams), sw=subwindow)
    maskpath = 'masks_{mt}_{sl}_{algparams_}/{c}_{s}_{mt}.mask'.format(**d)
    s = 'results_{m}_combinedDICOM/{c}_*_{s}/{c}_*_{s}_{p}'.format(**d)
    inpath = glob.glob(s)[0]
    outpath = 'rois_{mt}_{sl}_{algparams_}/{c}_x_x_{s}_{m}_{p}_{mt}.txt'.format(**d)
    args = [SELECT_VOXELS]
    #args += ['-v']
    if subwindow:
        args += ['-s %s' % subwindow_to_str(subwindow)]
    args += ['-m %s' % maskpath]
    args += ['-i "%s"' % inpath]
    args += ['-o "%s"' % outpath]
    cmd = ' '.join(args)
    return {
            'name': '{c}_{s}_{algparams_}'.format(**d),
            'actions': [(create_folder, [dirname(outpath)]),
                    cmd],
            'file_dep': [maskpath],
            'targets': [outpath],
            'clean': True,
            }

def task_select_roi_cancer():
    """Select cancer ROIs from the pmap DICOMs."""
    for sample in SAMPLES:
        for scan in sample['scans']:
            case = sample['case']
            masktype = 'cancer'
            yield get_task_select_roi(case, scan, 'Mono', 'ADCm', masktype,
                    ['default'])

def task_select_roi_auto():
    """Select automatic ROIs from the pmap DICOMs."""
    for sample in SAMPLES:
        for scan in sample['scans']:
            for algparams in itertools.product(*FIND_ROI_PARAMS):
                case = sample['case']
                masktype = 'auto'
                yield get_task_select_roi(case, scan, 'Mono', 'ADCm', masktype,
                        map(str, algparams))

def task_evaluate_autoroi():
    """Evaluate auto-ROI prediction ability by ROC AUC and correlation with
    Gleason score."""
    outfile = 'autoroi_evaluation.txt'
    d = dict(sl=SAMPLELIST, prg_auc=CALC_AUC, prg_corr=CORRELATION, o=outfile)
    cmds = ['echo `date` > {o}'.format(**d)]
    for algparams in itertools.product(*FIND_ROI_PARAMS):
        d['algparams_'] = '_'.join(map(str, algparams))
        d['i'] = 'rois_auto_{sl}_{algparams_}'.format(**d)
        s = 'echo -n {algparams_} >> {o}'
        cmds.append(s.format(**d))
        s = r'echo -n \\t`{prg_auc} -s patients.txt -l score -g 3+3 -m {i}/* -a --autoflip` >> {o}'
        cmds.append(s.format(**d))
        s = r'echo -n \\t`{prg_corr} -s patients.txt -l score -g 3+3 3+4 -a -m {i}/*` >> {o}'
        cmds.append(s.format(**d))
        s = r'echo \\t`{prg_corr} -s patients.txt -l score -a -m {i}/*` >> {o}'
        cmds.append(s.format(**d))
    return {
            'actions': cmds,
            'task_dep': ['select_roi_auto'],
            'targets': [outfile],
            'clean': True,
            }
<|endoftext|>


<commit_before>"""
Copyright (c) 2016 Stepan Fedorko-Bartos, Ceegan Hale

Under MIT License - https://github.com/Step7750/ScheduleStorm/blob/master/LICENSE.md

This file is a resource for Schedule Storm - https://github.com/Step7750/ScheduleStorm
"""

import threading
import requests
import pymongo
from bs4 import BeautifulSoup
import time
import re
from ldap3 import Server, Connection, SUBTREE, ALL, LEVEL
from queue import Queue
from .University import University

class UAlberta(University):
    def __init__(self, settings):
        super().__init__(settings)
        self.db = pymongo.MongoClient().ScheduleStorm

        self.db.UAlbertaProfessor.create_index([("uid", pymongo.ASCENDING)], unique=True)

    def parseCourseDescription(self, req):
        """
        Removes unnessary non-letters from the req

        :param req: **string** requisite form the description of a course
        :return: **string**
        """
        char = 1
        while not req[char].isalpha():
                        char += 1
        return req[char:]

    def scrapeCourseDesc(self, conn, termid):
        """
        Retrieves all course descriptions then parses the course requisites and notes then upserts for every entry in
        the query results

        :param conn: **ldap connection object**
        :param termid: **string/int** Term ID to get courses for
        :return: **string**
        """

        self.log.info('obtaining course descriptions')

        # Page queries course descriptions with the search base
        searchBase = 'term=' + termid + ', ou=calendar, dc=ualberta, dc=ca'
        entry_list = conn.extend.standard.paged_search(search_base=searchBase, search_filter='(course=*)',
                                                       search_scope=LEVEL,
                                                       attributes=['catalog', 'courseDescription', 'courseTitle',
                                                                   'subject', 'units'], paged_size=400, generator=False)

        # for entry in list, parse and upsert course descriptions
        for entry in entry_list:

            # initialize course description dict
            courseDesc = {
                'coursenum': entry['attributes']['catalog'],
                'subject': entry['attributes']['subject'],
                'name': entry['attributes']['courseTitle'],
                'units': entry['attributes']['units']
            }

            # Does the course have a description?
            if 'courseDescription' in entry['attributes']:
                desc = entry['attributes']['courseDescription']

                # Removes "See note (x) above" from description?
                if "See Note" in desc:
                    desc = desc.split("See Note", 1)[0]

                # Does the course have a prerequisite?
                if 'Prerequisite' in desc:

                    # Splits the prerequisite from the description
                    info = desc.split("Prerequisite", 1)
                    prereq = self.parseCourseDescription(info[1])
                    desc = info[0]

                    # Does prerequisite have a corequisite inside of it
                    if "Corequisite" in prereq or "corequisite" in prereq:

                        #Splits the corequisite from the prereq
                        if "Corequisite" in prereq:
                            info = prereq.split("Corequisite", 1)
                        elif "corequisite" in prereq:
                            info = prereq.split("corequisite", 1)

                        prereq = info[0]

                        # Removes any "and " leftover from the splitting
                        if prereq[-4:] == "and ":
                            prereq = prereq[:-4]

                        # if the coreq is different from the prereq
                        if len(info[1]) != 1:
                            corereq = self.parseCourseDescription(info[1])
                            if prereq == "or ":
                                prereq = corereq
                            else:
                                if corereq != prereq:
                                    courseDesc['coreq'] = corereq

                    # Splits the note form the prereq
                    if "Note:" in prereq:
                        note = prereq.split("Note:", 1)
                        courseDesc['notes'] = note[1]
                        prereq = note[0]

                    courseDesc['prereq'] = prereq

                # splits the antireq from the desc
                if "Antirequisite" in desc:
                    antireq = desc.split("Antirequisite", 1)[1]
                    antireq = self.parseCourseDescription(antireq)
                    courseDesc['antireq'] = antireq
                    desc = antireq[0]

                # removes leftover info from the desc split
                if desc[-4:] == "and ":
                            desc = desc[:-4]
                courseDesc['desc'] = desc

            # Upserts course description

            self.updateCourseDesc(courseDesc)

    def UidToName(self, uid):
        """
        Returns the name of the prof with the specified UID

        :param uid: **string** UID of the given prof
        :return: **string** Name of the prof if successful, UID if not
        """
        professor = self.db.UAlbertaProfessor.find({"uid": uid})
        if professor.count() == 0:
            # There must have been an issue when obtaining the data, just use the UID temporarily
            return uid
        else:
            # We got the name, return it
            professor = professor[0]['Name']
        return professor

    def scrapeCourseList(self, conn, termid):
        """
        Queries the course list with the termid, matches the professor to the course, upserts the initial dictionary
        then matches additional data to the object

        :param conn: **ldap connection object**
        :param termid: **string/int** Term ID to get courses for
        :return:
        """
        searchBase = 'term=' + termid + ', ou=calendar, dc=ualberta, dc=ca'
        entry_list = conn.extend.standard.paged_search(search_base=searchBase,
                                                       search_filter='(&(!(textbook=*))(class=*)(!(classtime=*)))',
                                                       search_scope=SUBTREE,
                                                       attributes=['asString', 'class', 'term', 'campus',
                                                                   'section', 'component', 'enrollStatus',
                                                                   'course', 'instructorUid'],
                                                       paged_size=400,
                                                       generator=False)

        # Searches for additional information
        times_list = conn.extend.standard.paged_search(search_base=searchBase,
                                                       search_filter='(&(!(textbook=*))(class=*)(classtime=*))',
                                                       search_scope=SUBTREE,
                                                       attributes=['day', 'class', 'startTime', 'endTime',
                                                                   'location'],
                                                       paged_size=400,
                                                       generator=False)

        # We want to scrape professor names from their UID's
        q = Queue()

        self.log.info("Filling up the Queue with Prof UIDs")

        # Fill queue with unique prof names
        queuedProfs = {}

        for entry in entry_list:
            # Ensure this class has teachers
            if 'instructorUid' in entry['attributes']:

                # We don't want to request duplicates
                if entry['attributes']['instructorUid'][0] not in queuedProfs:
                    q.put(entry['attributes']['instructorUid'][0])

                    # Add to the queuedProfs to avoid dupes
                    queuedProfs[entry['attributes']['instructorUid'][0]] = True

        # Start up the threads
        for i in range(self.settings["uidConcurrency"]):
            concurrentScraper = UIDScraper(q, self.db, self.log)
            concurrentScraper.daemon = True
            concurrentScraper.start()

        # Wait until the threads are done
        q.join()

        self.log.info('Parsing course data')

        # for each entry in list, upsert course into db
        for entry in entry_list:
            info = str(entry['attributes']['asString']).split(" ")

            # Seperates the subject from the coursenum
            if not info[1].isdigit():
                subject = info[0] + " " + info[1]
                coursenum = info[2]
            else:
                subject = info[0]
                coursenum = info[1]

            # Does the entry have an enrollStatus
            if entry['attributes']['enrollStatus'] == "O":
                status = "Open"
            elif entry['attributes']['enrollStatus'] == "C":
                status = "Closed"
            else:
                status = entry['attributes']['enrollStatus']

            # Initializes upsert dict
            courseList = {"subject": subject, "term": entry['attributes']['term'][0], "coursenum": coursenum,
                          "id": str(entry['attributes']['class']), "location": str(entry['attributes']['campus']),
                          "type": entry['attributes']['component'], "status": status,
                          'section': entry['attributes']['section'], "group": entry['attributes']['course'],
                          "times": ["N/A"], "rooms": ["N/A"]}

            # Does the entry have a instructor assigned to it
            if 'instructorUid' in entry['attributes']:
                courseList['teachers'] = [self.UidToName(entry['attributes']['instructorUid'][0])]
            else:
                courseList['teachers'] = ["N/A"]

            # for entry in list, match the days, startTime, endTime, and locations to course
            for entry_times in times_list:

                if entry_times['attributes']['class'] == courseList['id']:
                    # Combines day, startTime, endTime into a duration
                    duration = " "
                    duration = duration.join(
                        (entry_times['attributes']['day'][0], entry_times['attributes']['startTime'][0].replace(" ", ""),
                         entry_times['attributes']['endTime'][0].replace(" ", "")))

                    # Adds '-' btw the times
                    duration = re.sub(r'^((.*?\s.*?){1})\s', r'\1 - ', duration)
                    courseList['times'] = [duration]

                    # Does the class have an assigned classroom
                    if 'location' in entry['attributes']:
                        courseList['rooms'] = [entry['attributes']['location']]
                    times_list.remove(entry_times)
            # Upserts course into db
            self.updateClass(courseList)

    def scrapeTerms(self, conn):
        """
        Retrieves all course descriptions then parses the course requisites and notes then upserts for every entry in
        the query results

        :param conn: **ldap connection object**
        :return: **dict** has two keys term and termTitle, values are matched to their respective keys
        """

        # Page queries all terms
        conn.search(search_base='ou=calendar, dc=ualberta, dc=ca', search_filter='(term=*)', search_scope=LEVEL,
                    attributes=['term', 'termTitle'])
        terms = []

        # Gets the four most recent terms
        for item in range(1, 5):
            entry = conn.entries[len(conn.entries)-item]
            termDict = {"id": str(entry['term']), "name": str(entry['termTitle']).replace("Term ", "")}

            terms.append(termDict)

        # Adds term to term DB
        self.updateTerms(terms)

        # Returns current terms
        return terms

    def updateFaculties(self, conn):
        """
        Updates the faculties with the current terms as the search base

        :param conn: **ldap connection object**
        :return:
        """
        self.log.info("Getting faculty list")

        # Gets all recent terms and cycles through them
        for term in self.scrapeTerms(conn):

            # Sets the search base for the query
            searchBase = 'term='+term['id']+', ou=calendar, dc=ualberta, dc=ca'
            self.log.info("Updating faculties with search base " + searchBase)

            # Page queries all faculties in current term
            entry_list = conn.extend.standard.paged_search(search_base=searchBase,
                                                           search_filter='(term=*)',
                                                           search_scope=LEVEL,
                                                           attributes=['subject', 'subjectTitle', 'faculty'],
                                                           paged_size=400,
                                                           generator=False)

            subjectList = []
            # For each entry in list updates the faculty
            for entry in entry_list:
                if 'subject' in entry['attributes'] and entry['attributes']['subject'] not in subjectList:
                    subjectDict = {'subject': entry['attributes']['subject'],
                                   'faculty': entry['attributes']['faculty'],
                                   'name': entry['attributes']['subjectTitle']}
                    subjectList.append(subjectDict['subject'])

                    self.updateSubject(subjectDict)
        self.log.info('Finished updating faculties')

    def scrape(self):
        """
        Scraping thread that obtains updated course info

        :return:
        """
        # Establish connection to LDAP server
        server = Server('directory.srv.ualberta.ca', get_info=ALL)
        conn = Connection(server, auto_bind=True)

        # Updates faculties
        self.updateFaculties(conn)

        # Get list of current terms
        terms = self.getTerms()

        # For each term, get the courses
        for term in terms:
            self.log.info('Obtaining ' + terms[term] + ' course data with id ' + term)
            self.scrapeCourseList(conn, term)
            self.scrapeCourseDesc(conn, term)

        self.log.info('Finished scraping for UAlberta data')

class UIDScraper(threading.Thread):
    """
    Thread that gets UID's from the passed in queue and inserts the prof's data from UAlberta
    """

    def __init__(self, q, db, log):
        threading.Thread.__init__(self)
        self.q = q
        self.db = db
        self.log = log

    def run(self):
        """
        Scraping thread that gets a UID and inserts the returned prof data into the DB

        :return:
        """
        while not self.q.empty():
            # Get this UID from the queue
            thisuid = self.q.get()

            if thisuid:
                # Check if its already in the DB
                uidExists = self.db.UAlbertaProfessor.find({"uid": thisuid})

                if uidExists.count() == 0:
                    try:
                        # Get the prof data from the UAlberta directory
                        r = requests.get("http://webapps.srv.ualberta.ca/search/?type=simple&uid=true&c=" + thisuid,
                                        timeout=20)

                        # Check if the HTTP status code is ok
                        if r.status_code == requests.codes.ok:
                            # Parse the HTML
                            soup = BeautifulSoup(r.text, "lxml")
                            for tag in soup.find_all("b"):
                                info = tag.text
                                if info != "Dr " and info != "Prof ":
                                    professor = info
                                    break

                            self.log.info('Adding UID ' + thisuid + ' to UAlbertaProfessor db, Name: ' + professor)

                            # Upsert the data
                            self.db.UAlbertaProfessor.update({"uid": thisuid},
                                                             {'$set': {"uid": thisuid, "Name": professor}},
                                                             upsert=True)
                        else:
                            self.log.error("Improper HTTP Status for UID " + thisuid)
                    except:
                        self.log.error("Failed to obtain name for " + thisuid)

                # We're done with this class
                self.q.task_done()
            else:
                # No more items in the queue, stop the loop
                break
<commit_msg>Fixes graduate courses being in seperate faculties then ugrad courses<commit_after>"""
Copyright (c) 2016 Stepan Fedorko-Bartos, Ceegan Hale

Under MIT License - https://github.com/Step7750/ScheduleStorm/blob/master/LICENSE.md

This file is a resource for Schedule Storm - https://github.com/Step7750/ScheduleStorm
"""

import threading
import requests
import pymongo
from bs4 import BeautifulSoup
import time
import re
from ldap3 import Server, Connection, SUBTREE, ALL, LEVEL
from queue import Queue
from .University import University

class UAlberta(University):
    def __init__(self, settings):
        super().__init__(settings)
        self.db = pymongo.MongoClient().ScheduleStorm

        self.db.UAlbertaProfessor.create_index([("uid", pymongo.ASCENDING)], unique=True)

    def parseCourseDescription(self, req):
        """
        Removes unnessary non-letters from the req

        :param req: **string** requisite form the description of a course
        :return: **string**
        """
        char = 1
        while not req[char].isalpha():
                        char += 1
        return req[char:]

    def scrapeCourseDesc(self, conn, termid):
        """
        Retrieves all course descriptions then parses the course requisites and notes then upserts for every entry in
        the query results

        :param conn: **ldap connection object**
        :param termid: **string/int** Term ID to get courses for
        :return: **string**
        """

        self.log.info('obtaining course descriptions')

        # Page queries course descriptions with the search base
        searchBase = 'term=' + termid + ', ou=calendar, dc=ualberta, dc=ca'
        entry_list = conn.extend.standard.paged_search(search_base=searchBase, search_filter='(course=*)',
                                                       search_scope=LEVEL,
                                                       attributes=['catalog', 'courseDescription', 'courseTitle',
                                                                   'subject', 'units'], paged_size=400, generator=False)

        # for entry in list, parse and upsert course descriptions
        for entry in entry_list:

            # initialize course description dict
            courseDesc = {
                'coursenum': entry['attributes']['catalog'],
                'subject': entry['attributes']['subject'],
                'name': entry['attributes']['courseTitle'],
                'units': entry['attributes']['units']
            }

            # Does the course have a description?
            if 'courseDescription' in entry['attributes']:
                desc = entry['attributes']['courseDescription']

                # Removes "See note (x) above" from description?
                if "See Note" in desc:
                    desc = desc.split("See Note", 1)[0]

                # Does the course have a prerequisite?
                if 'Prerequisite' in desc:

                    # Splits the prerequisite from the description
                    info = desc.split("Prerequisite", 1)
                    prereq = self.parseCourseDescription(info[1])
                    desc = info[0]

                    # Does prerequisite have a corequisite inside of it
                    if "Corequisite" in prereq or "corequisite" in prereq:

                        #Splits the corequisite from the prereq
                        if "Corequisite" in prereq:
                            info = prereq.split("Corequisite", 1)
                        elif "corequisite" in prereq:
                            info = prereq.split("corequisite", 1)

                        prereq = info[0]

                        # Removes any "and " leftover from the splitting
                        if prereq[-4:] == "and ":
                            prereq = prereq[:-4]

                        # if the coreq is different from the prereq
                        if len(info[1]) != 1:
                            corereq = self.parseCourseDescription(info[1])
                            if prereq == "or ":
                                prereq = corereq
                            else:
                                if corereq != prereq:
                                    courseDesc['coreq'] = corereq

                    # Splits the note form the prereq
                    if "Note:" in prereq:
                        note = prereq.split("Note:", 1)
                        courseDesc['notes'] = note[1]
                        prereq = note[0]

                    courseDesc['prereq'] = prereq

                # splits the antireq from the desc
                if "Antirequisite" in desc:
                    antireq = desc.split("Antirequisite", 1)[1]
                    antireq = self.parseCourseDescription(antireq)
                    courseDesc['antireq'] = antireq
                    desc = antireq[0]

                # removes leftover info from the desc split
                if desc[-4:] == "and ":
                            desc = desc[:-4]
                courseDesc['desc'] = desc

            # Upserts course description

            self.updateCourseDesc(courseDesc)

    def UidToName(self, uid):
        """
        Returns the name of the prof with the specified UID

        :param uid: **string** UID of the given prof
        :return: **string** Name of the prof if successful, UID if not
        """
        professor = self.db.UAlbertaProfessor.find({"uid": uid})
        if professor.count() == 0:
            # There must have been an issue when obtaining the data, just use the UID temporarily
            return uid
        else:
            # We got the name, return it
            professor = professor[0]['Name']
        return professor

    def scrapeCourseList(self, conn, termid):
        """
        Queries the course list with the termid, matches the professor to the course, upserts the initial dictionary
        then matches additional data to the object

        :param conn: **ldap connection object**
        :param termid: **string/int** Term ID to get courses for
        :return:
        """
        searchBase = 'term=' + termid + ', ou=calendar, dc=ualberta, dc=ca'
        entry_list = conn.extend.standard.paged_search(search_base=searchBase,
                                                       search_filter='(&(!(textbook=*))(class=*)(!(classtime=*)))',
                                                       search_scope=SUBTREE,
                                                       attributes=['asString', 'class', 'term', 'campus',
                                                                   'section', 'component', 'enrollStatus',
                                                                   'course', 'instructorUid'],
                                                       paged_size=400,
                                                       generator=False)

        # Searches for additional information
        times_list = conn.extend.standard.paged_search(search_base=searchBase,
                                                       search_filter='(&(!(textbook=*))(class=*)(classtime=*))',
                                                       search_scope=SUBTREE,
                                                       attributes=['day', 'class', 'startTime', 'endTime',
                                                                   'location'],
                                                       paged_size=400,
                                                       generator=False)

        # We want to scrape professor names from their UID's
        q = Queue()

        self.log.info("Filling up the Queue with Prof UIDs")

        # Fill queue with unique prof names
        queuedProfs = {}

        for entry in entry_list:
            # Ensure this class has teachers
            if 'instructorUid' in entry['attributes']:

                # We don't want to request duplicates
                if entry['attributes']['instructorUid'][0] not in queuedProfs:
                    q.put(entry['attributes']['instructorUid'][0])

                    # Add to the queuedProfs to avoid dupes
                    queuedProfs[entry['attributes']['instructorUid'][0]] = True

        # Start up the threads
        for i in range(self.settings["uidConcurrency"]):
            concurrentScraper = UIDScraper(q, self.db, self.log)
            concurrentScraper.daemon = True
            concurrentScraper.start()

        # Wait until the threads are done
        q.join()

        self.log.info('Parsing course data')

        # for each entry in list, upsert course into db
        for entry in entry_list:
            info = str(entry['attributes']['asString']).split(" ")

            # Seperates the subject from the coursenum
            if not info[1].isdigit():
                subject = info[0] + " " + info[1]
                coursenum = info[2]
            else:
                subject = info[0]
                coursenum = info[1]

            # Does the entry have an enrollStatus
            if entry['attributes']['enrollStatus'] == "O":
                status = "Open"
            elif entry['attributes']['enrollStatus'] == "C":
                status = "Closed"
            else:
                status = entry['attributes']['enrollStatus']

            # Initializes upsert dict
            courseList = {"subject": subject, "term": entry['attributes']['term'][0], "coursenum": coursenum,
                          "id": str(entry['attributes']['class']), "location": str(entry['attributes']['campus']),
                          "type": entry['attributes']['component'], "status": status,
                          'section': entry['attributes']['section'], "group": entry['attributes']['course'],
                          "times": ["N/A"], "rooms": ["N/A"]}

            # Does the entry have a instructor assigned to it
            if 'instructorUid' in entry['attributes']:
                courseList['teachers'] = [self.UidToName(entry['attributes']['instructorUid'][0])]
            else:
                courseList['teachers'] = ["N/A"]

            # for entry in list, match the days, startTime, endTime, and locations to course
            for entry_times in times_list:

                if entry_times['attributes']['class'] == courseList['id']:
                    # Combines day, startTime, endTime into a duration
                    duration = " "
                    duration = duration.join(
                        (entry_times['attributes']['day'][0], entry_times['attributes']['startTime'][0].replace(" ", ""),
                         entry_times['attributes']['endTime'][0].replace(" ", "")))

                    # Adds '-' btw the times
                    duration = re.sub(r'^((.*?\s.*?){1})\s', r'\1 - ', duration)
                    courseList['times'] = [duration]

                    # Does the class have an assigned classroom
                    if 'location' in entry['attributes']:
                        courseList['rooms'] = [entry['attributes']['location']]
                    times_list.remove(entry_times)
            # Upserts course into db
            self.updateClass(courseList)

    def scrapeTerms(self, conn):
        """
        Retrieves all course descriptions then parses the course requisites and notes then upserts for every entry in
        the query results

        :param conn: **ldap connection object**
        :return: **dict** has two keys term and termTitle, values are matched to their respective keys
        """

        # Page queries all terms
        conn.search(search_base='ou=calendar, dc=ualberta, dc=ca', search_filter='(term=*)', search_scope=LEVEL,
                    attributes=['term', 'termTitle'])
        terms = []

        # Gets the seven most recent terms
        for item in range(1, 7):
            entry = conn.entries[len(conn.entries)-item]
            termDict = {"id": str(entry['term']), "name": str(entry['termTitle']).replace("Term ", "")}

            terms.append(termDict)

        # Adds term to term DB
        self.updateTerms(terms)

        # Returns current terms
        return terms

    def updateFaculties(self, conn):
        """
        Updates the faculties with the current terms as the search base

        :param conn: **ldap connection object**
        :return:
        """
        self.log.info("Getting faculty list")

        # Gets all recent terms and cycles through them
        for term in self.scrapeTerms(conn):

            # Sets the search base for the query
            searchBase = 'term='+term['id']+', ou=calendar, dc=ualberta, dc=ca'
            self.log.info("Updating faculties with search base " + searchBase)

            # Page queries all faculties in current term
            entry_list = conn.extend.standard.paged_search(search_base=searchBase,
                                                           search_filter='(term=*)',
                                                           search_scope=LEVEL,
                                                           attributes=['subject', 'subjectTitle', 'faculty', 'career'],
                                                           paged_size=400,
                                                           generator=False)

            ugrad = []
            # For each entry in list updates the faculty
            for entry in entry_list:
                if 'subject' in entry['attributes']:
                    subjectDict = {'subject': entry['attributes']['subject'],
                                   'faculty': entry['attributes']['faculty'],
                                   'name': entry['attributes']['subjectTitle']}
                    if entry['attributes']['career'] == 'UGRD':
                        ugrad.append(subjectDict['subject'])
                        self.updateSubject(subjectDict)
                    elif entry['attributes']['career'] == 'GRAD' and subjectDict['subject'] not in ugrad:
                        self.updateSubject(subjectDict)
        self.log.info('Finished updating faculties')

    def scrape(self):
        """
        Scraping thread that obtains updated course info

        :return:
        """
        # Establish connection to LDAP server
        server = Server('directory.srv.ualberta.ca', get_info=ALL)
        conn = Connection(server, auto_bind=True)

        # Updates faculties
        self.updateFaculties(conn)

        # Get list of current terms
        terms = self.getTerms()

        # For each term, get the courses
        for term in terms:
            self.log.info('Obtaining ' + terms[term] + ' course data with id ' + term)
            self.scrapeCourseList(conn, term)
            self.scrapeCourseDesc(conn, term)

        self.log.info('Finished scraping for UAlberta data')

class UIDScraper(threading.Thread):
    """
    Thread that gets UID's from the passed in queue and inserts the prof's data from UAlberta
    """

    def __init__(self, q, db, log):
        threading.Thread.__init__(self)
        self.q = q
        self.db = db
        self.log = log

    def run(self):
        """
        Scraping thread that gets a UID and inserts the returned prof data into the DB

        :return:
        """
        while not self.q.empty():
            # Get this UID from the queue
            thisuid = self.q.get()

            if thisuid:
                # Check if its already in the DB
                uidExists = self.db.UAlbertaProfessor.find({"uid": thisuid})

                if uidExists.count() == 0:
                    try:
                        # Get the prof data from the UAlberta directory
                        r = requests.get("http://webapps.srv.ualberta.ca/search/?type=simple&uid=true&c=" + thisuid,
                                        timeout=20)

                        # Check if the HTTP status code is ok
                        if r.status_code == requests.codes.ok:
                            # Parse the HTML
                            soup = BeautifulSoup(r.text, "lxml")
                            for tag in soup.find_all("b"):
                                info = tag.text
                                if info != "Dr " and info != "Prof ":
                                    professor = info
                                    break

                            self.log.info('Adding UID ' + thisuid + ' to UAlbertaProfessor db, Name: ' + professor)

                            # Upsert the data
                            self.db.UAlbertaProfessor.update({"uid": thisuid},
                                                             {'$set': {"uid": thisuid, "Name": professor}},
                                                             upsert=True)
                        else:
                            self.log.error("Improper HTTP Status for UID " + thisuid)
                    except:
                        self.log.error("Failed to obtain name for " + thisuid)

                # We're done with this class
                self.q.task_done()
            else:
                # No more items in the queue, stop the loop
                break
<|endoftext|>


<commit_before>import copy
import logging
from collections import defaultdict

log = logging.getLogger(__name__)

class Datastore(object):
    store = defaultdict(dict)
    on_change_cb = defaultdict(list)

    def __init__(self, domain):
        self.domain = domain

    def on_change(self, cb):
        Datastore.on_change_cb[self.domain].append(cb)
        return self

    def __setitem__(self, key, value):
        if key in Datastore.store[self.domain] and \
           Datastore.store[self.domain][key] == value:
            return self

        value = copy.deepcopy(value)

        Datastore.store[self.domain][key] = value
        for cb in Datastore.on_change_cb[self.domain]:
            cb(key, value)

        return self

    def __getitem__(self, key):
        return Datastore.store[self.domain][key]

    def update(self, d):
        """Similar to dict.update(). Call callbacks on each changed key."""
        for k, v in d.items():
            self[k] = v

    def init(self, d):
        """Only sets values that are not set already.
        No callbacks are called.
        """
        for k, v in d.items():
            if k not in Datastore.store[self.domain]:
                Datastore.store[self.domain][k] = copy.deepcopy(v)
<commit_msg>extend Datastore with __delitem__() and __contains__()<commit_after>import copy
import logging
from collections import defaultdict

log = logging.getLogger(__name__)

class Datastore(object):
    store = defaultdict(dict)
    on_change_cb = defaultdict(list)

    def __init__(self, domain):
        self.domain = domain

    def on_change(self, cb):
        Datastore.on_change_cb[self.domain].append(cb)
        return self

    def __setitem__(self, key, value):
        if key in Datastore.store[self.domain] and \
           Datastore.store[self.domain][key] == value:
            return self

        value = copy.deepcopy(value)

        Datastore.store[self.domain][key] = value
        for cb in Datastore.on_change_cb[self.domain]:
            cb(key, value)

        return self

    def __getitem__(self, key):
        return Datastore.store[self.domain][key]

    def __delitem__(self, key):
        del Datastore.store[self.domain][key]
        return self

    def __contains__(self, key):
        return key in Datastore.store[self.domain]

    def update(self, d):
        """Similar to dict.update(). Call callbacks on each changed key."""
        for k, v in d.items():
            self[k] = v

    def init(self, d):
        """Only sets values that are not set already.
        No callbacks are called.
        """
        for k, v in d.items():
            if k not in Datastore.store[self.domain]:
                Datastore.store[self.domain][k] = copy.deepcopy(v)
<|endoftext|>


<commit_before>#!/usr/bin/python2.4
#
# Copyright (c) 2004-2006 rPath, Inc.
#
# All rights reserved
#

import testsuite
testsuite.setup()
import simplejson

from mcp import client
from mcp import mcp_error

import mcp_helper

class ClientTest(mcp_helper.MCPTest):
    def queueResponse(self, res, error = False):
        self.client.response.inbound.insert(0, simplejson.dumps((error, res)))

    def getCommand(self):
        if self.client.command.connection.sent:
            res = self.client.command.connection.sent.pop()[1]
            self.failIf('action' not in res, "command is missing action")
            return simplejson.loads(res)
        self.fail("No command was in queue")

    def testCheckValue(self):
        command = {'action' : 'testCommand'}
        self.assertRaises(AssertionError, self.checkValue, command,
                          'missingKey', '')
        self.assertRaises(AssertionError, self.checkValue, command,
                          'action', 'wrong')
        assert self.checkValue(command, 'action', 'testCommand') is None

    def testBasicAttributes(self):
        assert self.client.command.connectionName == '/queue/test/command'
        assert self.client.response.connectionName == \
            '/topic/test/' + self.client.uuid

        res = self.client.slaveStatus()
        assert self.client.command.connection.sent[0][0] == \
            '/queue/test/command'

    def testSubmitBuild(self):
        build = self.getJsonBuild()

        self.client.submitJob(build)
        assert self.client.command.connection.sent[0][0] == \
            '/queue/test/command'

        command = simplejson.loads(self.client.command.connection.sent[0][1])

        assert command['data'] == build
        assert self.client.response.connectionName == \
            '/topic/test/' + command['uuid']

        self.checkValue(command, 'action', 'submitJob')
        self.checkValue(command, 'protocolVersion', 1)

    def testSlaveStatus(self):
        self.queueResponse({})
        res = self.client.slaveStatus()
        assert res == {}

        command = self.getCommand()
        self.checkValue(command, 'action', 'slaveStatus')

    def testStopJob(self):
        self.queueResponse(None)
        self.client.stopJob('dummy-jobId')

        command = self.getCommand()
        self.checkValue(command, 'action', 'stopJob')
        self.checkValue(command, 'jobId', 'dummy-jobId')

    def testJobStatus(self):
        self.queueResponse(('failed', 'for no reason'))
        res = self.client.jobStatus('test-job')
        assert res == ['failed', 'for no reason']

        command = self.getCommand()
        self.checkValue(command, 'action', 'jobStatus')
        self.checkValue(command, 'jobId', 'test-job')

    def testStopSlave(self):
        self.queueResponse(None)
        self.client.stopSlave('dummy-slaveId')

        command = self.getCommand()
        self.checkValue(command, 'action', 'stopSlave')
        self.checkValue(command, 'slaveId', 'dummy-slaveId')
        self.checkValue(command, 'delayed', True)

        self.queueResponse(None)
        self.client.stopSlave('dummy-slaveId', delayed = False)
        command = self.getCommand()
        self.checkValue(command, 'delayed', False)

    def testStopMaster(self):
        self.queueResponse(None)
        self.client.stopMaster('dummy-masterId')
        command = self.getCommand()
        self.checkValue(command, 'action', 'stopMaster')
        self.checkValue(command, 'masterId', 'dummy-masterId')
        self.checkValue(command, 'delayed', True)

        self.queueResponse(None)
        self.client.stopMaster('dummy-masterId', delayed = False)
        command = self.getCommand()
        self.checkValue(command, 'delayed', False)

    def testSetSlaveTTL(self):
        self.queueResponse(None)
        self.client.setSlaveTTL('dummy-slaveId', 60)
        command = self.getCommand()
        self.checkValue(command, 'action', 'setSlaveTTL')
        self.checkValue(command, 'slaveId', 'dummy-slaveId')
        self.checkValue(command, 'TTL', 60)

    def testSetSlaveLimit(self):
        self.queueResponse(None)
        self.client.setSlaveLimit('dummy-masterId', 3)
        command = self.getCommand()

        self.checkValue(command, 'action', 'setSlaveLimit')
        self.checkValue(command, 'masterId', 'dummy-masterId')
        self.checkValue(command, 'limit', 3)

if __name__ == "__main__":
    testsuite.main()
<commit_msg>client coverage<commit_after>#!/usr/bin/python2.4
#
# Copyright (c) 2004-2006 rPath, Inc.
#
# All rights reserved
#

import testsuite
testsuite.setup()
import simplejson

from mcp import client
from mcp import mcp_error

import mcp_helper

class ClientTest(mcp_helper.MCPTest):
    def queueResponse(self, res, error = False):
        self.client.response.inbound.insert(0, simplejson.dumps((error, res)))

    def getCommand(self):
        if self.client.command.connection.sent:
            res = self.client.command.connection.sent.pop()[1]
            self.failIf('action' not in res, "command is missing action")
            return simplejson.loads(res)
        self.fail("No command was in queue")

    def testCheckValue(self):
        command = {'action' : 'testCommand'}
        self.assertRaises(AssertionError, self.checkValue, command,
                          'missingKey', '')
        self.assertRaises(AssertionError, self.checkValue, command,
                          'action', 'wrong')
        assert self.checkValue(command, 'action', 'testCommand') is None

    def testBasicAttributes(self):
        assert self.client.command.connectionName == '/queue/test/command'
        assert self.client.response.connectionName == \
            '/topic/test/' + self.client.uuid

        res = self.client.slaveStatus()
        assert self.client.command.connection.sent[0][0] == \
            '/queue/test/command'

    def testSubmitBuild(self):
        build = self.getJsonBuild()

        self.client.submitJob(build)
        assert self.client.command.connection.sent[0][0] == \
            '/queue/test/command'

        command = simplejson.loads(self.client.command.connection.sent[0][1])

        assert command['data'] == build
        assert self.client.response.connectionName == \
            '/topic/test/' + command['uuid']

        self.checkValue(command, 'action', 'submitJob')
        self.checkValue(command, 'protocolVersion', 1)

    def testSlaveStatus(self):
        self.queueResponse({})
        res = self.client.slaveStatus()
        assert res == {}

        command = self.getCommand()
        self.checkValue(command, 'action', 'slaveStatus')

    def testStopJob(self):
        self.queueResponse(None)
        self.client.stopJob('dummy-jobId')

        command = self.getCommand()
        self.checkValue(command, 'action', 'stopJob')
        self.checkValue(command, 'jobId', 'dummy-jobId')

    def testJobStatus(self):
        self.queueResponse(('failed', 'for no reason'))
        res = self.client.jobStatus('test-job')
        assert res == ['failed', 'for no reason']

        command = self.getCommand()
        self.checkValue(command, 'action', 'jobStatus')
        self.checkValue(command, 'jobId', 'test-job')

    def testStopSlave(self):
        self.queueResponse(None)
        self.client.stopSlave('dummy-slaveId')

        command = self.getCommand()
        self.checkValue(command, 'action', 'stopSlave')
        self.checkValue(command, 'slaveId', 'dummy-slaveId')
        self.checkValue(command, 'delayed', True)

        self.queueResponse(None)
        self.client.stopSlave('dummy-slaveId', delayed = False)
        command = self.getCommand()
        self.checkValue(command, 'delayed', False)

    def testStopMaster(self):
        self.queueResponse(None)
        self.client.stopMaster('dummy-masterId')
        command = self.getCommand()
        self.checkValue(command, 'action', 'stopMaster')
        self.checkValue(command, 'masterId', 'dummy-masterId')
        self.checkValue(command, 'delayed', True)

        self.queueResponse(None)
        self.client.stopMaster('dummy-masterId', delayed = False)
        command = self.getCommand()
        self.checkValue(command, 'delayed', False)

    def testSetSlaveTTL(self):
        self.queueResponse(None)
        self.client.setSlaveTTL('dummy-slaveId', 60)
        command = self.getCommand()
        self.checkValue(command, 'action', 'setSlaveTTL')
        self.checkValue(command, 'slaveId', 'dummy-slaveId')
        self.checkValue(command, 'TTL', 60)

    def testSetSlaveLimit(self):
        self.queueResponse(None)
        self.client.setSlaveLimit('dummy-masterId', 3)
        command = self.getCommand()

        self.checkValue(command, 'action', 'setSlaveLimit')
        self.checkValue(command, 'masterId', 'dummy-masterId')
        self.checkValue(command, 'limit', 3)

    def testBrokenSend(self):
        self.assertRaises(AssertionError, self.client._send, test = 'test')
        def __del__(self, garbage):
            self._send(garbage = garbage)
        self.client.__del__ = __del__
        self.assertRaises(mcp_error.ProtocolError,
                          self.client.__del__, self.client, 'broken')

    def testMarshallError(self):
        self.queueResponse(('ProtocolError', 'just a test'), error = True)
        self.assertRaises(mcp_error.ProtocolError, self.client.jobStatus, '')
        self.queueResponse(('AssertionError', 'just a test'), error = True)
        self.assertRaises(Exception, self.client.jobStatus, '')

if __name__ == "__main__":
    testsuite.main()
<|endoftext|>


<commit_before>"""PyDoIt file for automating tasks."""

import glob
import os

from doit import get_var
#from doit.tools import Interactive
from doit.tools import create_folder
#from doit.tools import result_dep

import dwi.util

"""
Backends:
dbm: (default) It uses python dbm module.
json: Plain text using a json structure, it is slow but good for debugging.
sqlite3: (experimental) very slow implementation, support concurrent access.
"""

DOIT_CONFIG = {
    'backend': 'sqlite3',
    'default_tasks': [],
    'verbosity': 2,
    }

DWILIB = '~/src/dwilib/tools'
PMAP = DWILIB+'/pmap.py'
ANON = DWILIB+'/anonymize_dicom.py'
FIND_ROI = DWILIB+'/find_roi.py'
COMPARE_MASKS = DWILIB+'/compare_masks.py'
SELECT_VOXELS = DWILIB+'/select_voxels.py'
MODELS = 'Si SiN Mono MonoN Kurt KurtN Stretched StretchedN '\
        'Biexp BiexpN'.split()

SAMPLELIST = get_var('samplelist', 'all') # Sample list (train, test, etc)
SAMPLELIST_FILE = 'samples_%s.txt' % SAMPLELIST
SAMPLES = dwi.util.read_sample_list(SAMPLELIST_FILE)
SUBWINDOWS = dwi.util.read_subwindows('subwindows.txt')

# Common functions

def subwindow_to_str(subwindow):
    return ' '.join(map(str, subwindow))

def fit_cmd(model, subwindow, infiles, outfile):
    d = dict(prg=PMAP,
            m=model,
            sw=subwindow_to_str(subwindow),
            i=' '.join(infiles),
            o=outfile,
            )
    s = '{prg} -m {m} -s {sw} -d {i} -o {o}'.format(**d)
    return s

# Tasks

#def task_anonymize():
#    """Anonymize imaging data."""
#    files = glob.glob('dicoms/*/DICOMDIR') + glob.glob('dicoms/*/DICOM/*')
#    files.sort()
#    for f in files:
#        cmd = '{prg} -v -i {f}'.format(prg=ANON, f=f)
#        yield {
#                'name': f,
#                'actions': [cmd],
#                'file_dep': [f],
#                }

def task_fit():
    """Fit models to imaging data."""
    for case, scan in SUBWINDOWS.keys():
        subwindow = SUBWINDOWS[(case, scan)]
        outdir = 'pmaps'
        d = dict(c=case, s=scan, od=outdir)
        infiles = sorted(glob.glob('dicoms/{c}_*_hB_{s}*/DICOM/*'.format(**d)))
        if len(infiles) == 0:
            continue
        for model in MODELS:
            d['m'] = model
            outfile = '{od}/pmap_{c}_{s}_{m}.txt'.format(**d)
            cmd = fit_cmd(model, subwindow, infiles, outfile)
            yield {
                    'name': '{c}_{s}_{m}'.format(**d),
                    'actions': [(create_folder, [outdir]),
                            cmd],
                    'file_dep': infiles,
                    'targets': [outfile],
                    'clean': True,
                    }

def task_find_roi():
    """Find ROIs."""
    for sample in SAMPLES:
        case = sample['case']
        for scan in sample['scans']:
            d = dict(prg=FIND_ROI, sl=SAMPLELIST_FILE, c=case, s=scan)
            d['outdir_mask'] = 'masks_auto'
            d['outdir_fig'] = 'find_roi_images'
            d['outfile_mask'] = '{outdir_mask}/{c}_{s}_auto.mask'.format(**d)
            d['outfile_fig'] = '{outdir_fig}/autoroi_{c}_{s}.png'.format(**d)
            #infile = 'pmaps/pmap_{c}_{s}_MonoN.txt'.format(**d)
            #inmask = glob.glob('masks/{c}_{s}_1_*.mask'.format(**d))[0]
            cmd = '{prg} --samplelist {sl} --cases {c} --scans {s}'.format(**d)
            yield {
                    'name': '{c}_{s}'.format(**d),
                    'actions': [(create_folder, [d['outdir_mask']]),
                                (create_folder, [d['outdir_fig']]),
                            cmd],
                    #'file_dep': [infile],
                    'targets': [d['outfile_mask'], d['outfile_fig']],
                    'clean': True,
                    }

def task_compare_masks():
    """Compare ROI masks."""
    for case, scan in SUBWINDOWS.keys():
        subwindow = SUBWINDOWS[(case, scan)]
        d = dict(prg=COMPARE_MASKS,
                c=case,
                s=scan,
                w=subwindow_to_str(subwindow),
                )
        #file1 = 'masks/{c}_{s}_ca.mask'.format(**d)
        file1 = glob.glob('masks/{c}_{s}_1_*.mask'.format(**d))[0]
        file2 = 'masks_auto/{c}_{s}_auto.mask'.format(**d)
        outdir = 'roi_comparison'
        outfile = os.path.join(outdir, '{c}_{s}.txt'.format(**d))
        d['f1'] = file1
        d['f2'] = file2
        d['o'] = outfile
        cmd = '{prg} -s {w} {f1} {f2} > {o}'.format(**d)
        if not os.path.exists(file1):
            print 'Missing mask file %s' % file1
            continue
        yield {
                'name': '{c}_{s}'.format(**d),
                'actions': [(create_folder, [outdir]),
                        cmd],
                'file_dep': [file1, file2],
                'targets': [outfile],
                'clean': True,
                }

def get_task_select_roi(case, scan, model, param, subwindow, mask):
    # TODO: Cannot depend on mask files, as they are named irregularly.
    d = dict(c=case, s=scan, m=model, p=param, subwindow=subwindow, mask=mask)
    outdir = 'rois_{mask}'.format(**d)
    s = os.path.join('masks_{mask}', '{c}_{s}_*.mask')
    #maskpath = glob.glob(s.format(**d))[0]
    maskpath = s.format(**d)
    s = os.path.join('results_{m}_combinedDICOM', '{c}_*_{s}',
            '{c}_*_{s}_{p}')
    inpath = glob.glob(s.format(**d))[0]
    s = os.path.join(outdir, '{c}_x_x_{s}_{m}_{p}_{mask}.txt')
    outpath = s.format(**d)
    args = [SELECT_VOXELS]
    #args += ['-v']
    if subwindow:
        args += ['-s %s' % subwindow_to_str(subwindow)]
    args += ['-m %s' % maskpath]
    args += ['-i "%s"' % inpath]
    args += ['-o "%s"' % outpath]
    cmd = ' '.join(args)
    return {
            'name': '{c}_{s}'.format(**d),
            'actions': [(create_folder, [outdir]),
                    cmd],
            #'file_dep': [maskpath],
            'file_dep': [],
            'targets': [outpath],
            'clean': True,
            }

def task_select_roi_cancer():
    """Select cancer ROIs from the pmap DICOMs."""
    for sample in SAMPLES:
        for scan in sample['scans']:
            case = sample['case']
            subwin = None
            mask = 'cancer'
            yield get_task_select_roi(case, scan, 'Mono', 'ADCm', subwin, mask)

def task_select_roi_auto():
    """Select automatic ROIs from the pmap DICOMs."""
    for sample in SAMPLES:
        for scan in sample['scans']:
            case = sample['case']
            subwin = None
            mask = 'auto'
            yield get_task_select_roi(case, scan, 'Mono', 'ADCm', subwin, mask)
<commit_msg>Mask file dependence for roi selection task<commit_after>"""PyDoIt file for automating tasks."""

import glob
import os

from doit import get_var
#from doit.tools import Interactive
from doit.tools import create_folder
#from doit.tools import result_dep

import dwi.util

"""
Backends:
dbm: (default) It uses python dbm module.
json: Plain text using a json structure, it is slow but good for debugging.
sqlite3: (experimental) very slow implementation, support concurrent access.
"""

DOIT_CONFIG = {
    'backend': 'sqlite3',
    'default_tasks': [],
    'verbosity': 2,
    }

DWILIB = '~/src/dwilib/tools'
PMAP = DWILIB+'/pmap.py'
ANON = DWILIB+'/anonymize_dicom.py'
FIND_ROI = DWILIB+'/find_roi.py'
COMPARE_MASKS = DWILIB+'/compare_masks.py'
SELECT_VOXELS = DWILIB+'/select_voxels.py'
MODELS = 'Si SiN Mono MonoN Kurt KurtN Stretched StretchedN '\
        'Biexp BiexpN'.split()

SAMPLELIST = get_var('samplelist', 'all') # Sample list (train, test, etc)
SAMPLELIST_FILE = 'samples_%s.txt' % SAMPLELIST
SAMPLES = dwi.util.read_sample_list(SAMPLELIST_FILE)
SUBWINDOWS = dwi.util.read_subwindows('subwindows.txt')

# Common functions

def subwindow_to_str(subwindow):
    return ' '.join(map(str, subwindow))

def fit_cmd(model, subwindow, infiles, outfile):
    d = dict(prg=PMAP,
            m=model,
            sw=subwindow_to_str(subwindow),
            i=' '.join(infiles),
            o=outfile,
            )
    s = '{prg} -m {m} -s {sw} -d {i} -o {o}'.format(**d)
    return s

# Tasks

#def task_anonymize():
#    """Anonymize imaging data."""
#    files = glob.glob('dicoms/*/DICOMDIR') + glob.glob('dicoms/*/DICOM/*')
#    files.sort()
#    for f in files:
#        cmd = '{prg} -v -i {f}'.format(prg=ANON, f=f)
#        yield {
#                'name': f,
#                'actions': [cmd],
#                'file_dep': [f],
#                }

def task_fit():
    """Fit models to imaging data."""
    for case, scan in SUBWINDOWS.keys():
        subwindow = SUBWINDOWS[(case, scan)]
        outdir = 'pmaps'
        d = dict(c=case, s=scan, od=outdir)
        infiles = sorted(glob.glob('dicoms/{c}_*_hB_{s}*/DICOM/*'.format(**d)))
        if len(infiles) == 0:
            continue
        for model in MODELS:
            d['m'] = model
            outfile = '{od}/pmap_{c}_{s}_{m}.txt'.format(**d)
            cmd = fit_cmd(model, subwindow, infiles, outfile)
            yield {
                    'name': '{c}_{s}_{m}'.format(**d),
                    'actions': [(create_folder, [outdir]),
                            cmd],
                    'file_dep': infiles,
                    'targets': [outfile],
                    'clean': True,
                    }

def task_find_roi():
    """Find ROIs."""
    for sample in SAMPLES:
        case = sample['case']
        for scan in sample['scans']:
            d = dict(prg=FIND_ROI, sl=SAMPLELIST_FILE, c=case, s=scan)
            d['outdir_mask'] = 'masks_auto'
            d['outdir_fig'] = 'find_roi_images'
            d['outfile_mask'] = '{outdir_mask}/{c}_{s}_auto.mask'.format(**d)
            d['outfile_fig'] = '{outdir_fig}/autoroi_{c}_{s}.png'.format(**d)
            #infile = 'pmaps/pmap_{c}_{s}_MonoN.txt'.format(**d)
            #inmask = glob.glob('masks/{c}_{s}_1_*.mask'.format(**d))[0]
            cmd = '{prg} --samplelist {sl} --cases {c} --scans {s}'.format(**d)
            yield {
                    'name': '{c}_{s}'.format(**d),
                    'actions': [(create_folder, [d['outdir_mask']]),
                                (create_folder, [d['outdir_fig']]),
                            cmd],
                    #'file_dep': [infile],
                    'targets': [d['outfile_mask'], d['outfile_fig']],
                    'clean': True,
                    }

def task_compare_masks():
    """Compare ROI masks."""
    for case, scan in SUBWINDOWS.keys():
        subwindow = SUBWINDOWS[(case, scan)]
        d = dict(prg=COMPARE_MASKS,
                c=case,
                s=scan,
                w=subwindow_to_str(subwindow),
                )
        #file1 = 'masks/{c}_{s}_ca.mask'.format(**d)
        file1 = glob.glob('masks/{c}_{s}_1_*.mask'.format(**d))[0]
        file2 = 'masks_auto/{c}_{s}_auto.mask'.format(**d)
        outdir = 'roi_comparison'
        outfile = os.path.join(outdir, '{c}_{s}.txt'.format(**d))
        d['f1'] = file1
        d['f2'] = file2
        d['o'] = outfile
        cmd = '{prg} -s {w} {f1} {f2} > {o}'.format(**d)
        if not os.path.exists(file1):
            print 'Missing mask file %s' % file1
            continue
        yield {
                'name': '{c}_{s}'.format(**d),
                'actions': [(create_folder, [outdir]),
                        cmd],
                'file_dep': [file1, file2],
                'targets': [outfile],
                'clean': True,
                }

def get_task_select_roi(case, scan, model, param, subwindow, mask):
    # TODO: Cannot depend on mask files, as they are named irregularly.
    d = dict(c=case, s=scan, m=model, p=param, subwindow=subwindow, mask=mask)
    outdir = 'rois_{mask}'.format(**d)
    s = os.path.join('masks_{mask}', '{c}_{s}_{mask}.mask')
    maskpath = s.format(**d)
    s = os.path.join('results_{m}_combinedDICOM', '{c}_*_{s}',
            '{c}_*_{s}_{p}')
    inpath = glob.glob(s.format(**d))[0]
    s = os.path.join(outdir, '{c}_x_x_{s}_{m}_{p}_{mask}.txt')
    outpath = s.format(**d)
    args = [SELECT_VOXELS]
    #args += ['-v']
    if subwindow:
        args += ['-s %s' % subwindow_to_str(subwindow)]
    args += ['-m %s' % maskpath]
    args += ['-i "%s"' % inpath]
    args += ['-o "%s"' % outpath]
    cmd = ' '.join(args)
    return {
            'name': '{c}_{s}'.format(**d),
            'actions': [(create_folder, [outdir]),
                    cmd],
            'file_dep': [maskpath],
            'targets': [outpath],
            'clean': True,
            }

def task_select_roi_cancer():
    """Select cancer ROIs from the pmap DICOMs."""
    for sample in SAMPLES:
        for scan in sample['scans']:
            case = sample['case']
            subwin = None
            mask = 'cancer'
            yield get_task_select_roi(case, scan, 'Mono', 'ADCm', subwin, mask)

def task_select_roi_auto():
    """Select automatic ROIs from the pmap DICOMs."""
    for sample in SAMPLES:
        for scan in sample['scans']:
            case = sample['case']
            subwin = None
            mask = 'auto'
            yield get_task_select_roi(case, scan, 'Mono', 'ADCm', subwin, mask)
<|endoftext|>


<commit_before>from collections import defaultdict, namedtuple
import re
import inspect

BOOLEAN_FIELDS = (
    "search.regex", "searchable", "orderable", "regex"
)

DataColumn = namedtuple("DataColumn", ("name", "model_name", "filter"))

class DataTablesError(ValueError):
    pass

class DataTable(object):
    def __init__(self, params, model, query, columns):
        self.params = params
        self.model = model
        self.query = query
        self.data = {}
        self.columns = []
        self.columns_dict = {}
        self.search_func = lambda qs, s: qs

        for col in columns:
            name, model_name, filter_func = None, None, None

            if isinstance(col, DataColumn):
                self.columns.append(col)
                continue
            elif isinstance(col, tuple):
                # col is either 1. (name, model_name), 2. (name, filter) or 3. (name, model_name, filter)
                if len(col) == 3:
                    name, model_name, filter_func = col
                elif len(col) == 2:
                    # Work out the second argument. If it is a function then it's type 2, else it is type 1.
                    if callable(col[1]):
                        name, filter_func = col
                        model_name = name
                    else:
                        name, model_name = col
                else:
                    raise ValueError("Columns must be a tuple of 2 to 3 elements")
            else:
                # It's just a string
                name, model_name = col, col

            d = DataColumn(name=name, model_name=model_name, filter=filter_func)
            self.columns.append(d)
            self.columns_dict[d.name] = d

        for column in (col for col in self.columns if "." in col.model_name):
            self.query = self.query.join(column.model_name.split(".")[0], aliased=True)

    def query_into_dict(self, key_start):
        returner = defaultdict(dict)

        # Matches columns[number][key] with an [optional_value] on the end
        pattern = "{}(?:\[(\d+)\])?\[(\w+)\](?:\[(\w+)\])?".format(key_start)

        columns = (param for param in self.params if re.match(pattern, param))

        for param in columns:

            column_id, key, optional_subkey = re.search(pattern, param).groups()

            if column_id is None:
                returner[key] = self.coerce_value(key, self.params[param])
            elif optional_subkey is None:
                returner[int(column_id)][key] = self.coerce_value(key, self.params[param])
            else:
                # Oh baby a triple
                subdict = returner[int(column_id)].setdefault(key, {})
                subdict[optional_subkey] = self.coerce_value("{}.{}".format(key, optional_subkey),
                                                             self.params[param])

        return dict(returner)

    @staticmethod
    def coerce_value(key, value):
        try:
            return int(value)
        except ValueError:
            if key in BOOLEAN_FIELDS:
                return value == "true"

        return value

    def get_integer_param(self, param_name):
        if param_name not in self.params:
            raise DataTablesError("Parameter {} is missing".format(param_name))

        try:
            return int(self.params[param_name])
        except ValueError:
            raise DataTablesError("Parameter {} is invalid".format(param_name))

    def add_data(self, **kwargs):
        self.data.update(**kwargs)

    def json(self):
        try:
            return self._json()
        except DataTablesError as e:
            return {
                "error": str(e)
            }

    def get_column(self, column):
        if "." in column.model_name:
            column_path = column.model_name.split(".")
            relationship = getattr(self.model, column_path[0])
            model_column = getattr(relationship.property.mapper.entity, column_path[1])
        else:
            model_column = getattr(self.model, column.model_name)

        return model_column

    def searchable(self, func):
        self.search_func = func

    def _json(self):
        draw = self.get_integer_param("draw")
        start = self.get_integer_param("start")
        length = self.get_integer_param("length")

        columns = self.query_into_dict("columns")
        ordering = self.query_into_dict("order")
        search = self.query_into_dict("search")

        query = self.query
        total_records = query.count()

        if callable(self.search_func) and search.get("value", None):
            query = self.search_func(query, str(search["value"]))

        for order in ordering.values():
            direction, column = order["dir"], order["column"]

            if column not in columns:
                raise DataTablesError("Cannot order {}: column not found".format(column))

            if not columns[column]["orderable"]:
                continue

            column_name = columns[column]["data"]
            column = self.columns_dict[column_name]

            model_column = self.get_column(column)

            if isinstance(model_column, property):
                raise DataTablesError("Cannot order by column {} as it is a property".format(column.model_name))

            query = query.order_by(model_column.desc() if direction == "desc" else model_column.asc())

        filtered_records = query.count()

        if length > 0:
            query = query.slice(start, start + length)

        return {
            "draw": draw,
            "recordsTotal": total_records,
            "recordsFiltered": filtered_records,
            "data": [
                self.output_instance(instance) for instance in query.all()
            ]
        }

    def output_instance(self, instance):
        returner = {
            key.name: self.get_value(key, instance) for key in self.columns
        }

        if self.data:
            returner["DT_RowData"] = {
                k: v(instance) for k, v in self.data.items()
            }

        return returner

    def get_value(self, key, instance):
        attr = key.model_name
        if "." in attr:
            tmp_list=attr.split(".")
            attr=tmp_list[-1]
            for sub in tmp_list[:-1]:
                instance = getattr(instance, sub)

        if key.filter is not None:
            r = key.filter(instance)
        else:
            r = getattr(instance, attr)

        return r() if inspect.isroutine(r) else r
<commit_msg>Unicode search fix.<commit_after>from collections import defaultdict, namedtuple
import re
import inspect

BOOLEAN_FIELDS = (
    "search.regex", "searchable", "orderable", "regex"
)

DataColumn = namedtuple("DataColumn", ("name", "model_name", "filter"))

class DataTablesError(ValueError):
    pass

class DataTable(object):
    def __init__(self, params, model, query, columns):
        self.params = params
        self.model = model
        self.query = query
        self.data = {}
        self.columns = []
        self.columns_dict = {}
        self.search_func = lambda qs, s: qs

        for col in columns:
            name, model_name, filter_func = None, None, None

            if isinstance(col, DataColumn):
                self.columns.append(col)
                continue
            elif isinstance(col, tuple):
                # col is either 1. (name, model_name), 2. (name, filter) or 3. (name, model_name, filter)
                if len(col) == 3:
                    name, model_name, filter_func = col
                elif len(col) == 2:
                    # Work out the second argument. If it is a function then it's type 2, else it is type 1.
                    if callable(col[1]):
                        name, filter_func = col
                        model_name = name
                    else:
                        name, model_name = col
                else:
                    raise ValueError("Columns must be a tuple of 2 to 3 elements")
            else:
                # It's just a string
                name, model_name = col, col

            d = DataColumn(name=name, model_name=model_name, filter=filter_func)
            self.columns.append(d)
            self.columns_dict[d.name] = d

        for column in (col for col in self.columns if "." in col.model_name):
            self.query = self.query.join(column.model_name.split(".")[0], aliased=True)

    def query_into_dict(self, key_start):
        returner = defaultdict(dict)

        # Matches columns[number][key] with an [optional_value] on the end
        pattern = "{}(?:\[(\d+)\])?\[(\w+)\](?:\[(\w+)\])?".format(key_start)

        columns = (param for param in self.params if re.match(pattern, param))

        for param in columns:

            column_id, key, optional_subkey = re.search(pattern, param).groups()

            if column_id is None:
                returner[key] = self.coerce_value(key, self.params[param])
            elif optional_subkey is None:
                returner[int(column_id)][key] = self.coerce_value(key, self.params[param])
            else:
                # Oh baby a triple
                subdict = returner[int(column_id)].setdefault(key, {})
                subdict[optional_subkey] = self.coerce_value("{}.{}".format(key, optional_subkey),
                                                             self.params[param])

        return dict(returner)

    @staticmethod
    def coerce_value(key, value):
        try:
            return int(value)
        except ValueError:
            if key in BOOLEAN_FIELDS:
                return value == "true"

        return value

    def get_integer_param(self, param_name):
        if param_name not in self.params:
            raise DataTablesError("Parameter {} is missing".format(param_name))

        try:
            return int(self.params[param_name])
        except ValueError:
            raise DataTablesError("Parameter {} is invalid".format(param_name))

    def add_data(self, **kwargs):
        self.data.update(**kwargs)

    def json(self):
        try:
            return self._json()
        except DataTablesError as e:
            return {
                "error": str(e)
            }

    def get_column(self, column):
        if "." in column.model_name:
            column_path = column.model_name.split(".")
            relationship = getattr(self.model, column_path[0])
            model_column = getattr(relationship.property.mapper.entity, column_path[1])
        else:
            model_column = getattr(self.model, column.model_name)

        return model_column

    def searchable(self, func):
        self.search_func = func

    def _json(self):
        draw = self.get_integer_param("draw")
        start = self.get_integer_param("start")
        length = self.get_integer_param("length")

        columns = self.query_into_dict("columns")
        ordering = self.query_into_dict("order")
        search = self.query_into_dict("search")

        query = self.query
        total_records = query.count()

        if callable(self.search_func) and search.get("value", None):
            query = self.search_func(query, search["value"])

        for order in ordering.values():
            direction, column = order["dir"], order["column"]

            if column not in columns:
                raise DataTablesError("Cannot order {}: column not found".format(column))

            if not columns[column]["orderable"]:
                continue

            column_name = columns[column]["data"]
            column = self.columns_dict[column_name]

            model_column = self.get_column(column)

            if isinstance(model_column, property):
                raise DataTablesError("Cannot order by column {} as it is a property".format(column.model_name))

            query = query.order_by(model_column.desc() if direction == "desc" else model_column.asc())

        filtered_records = query.count()

        if length > 0:
            query = query.slice(start, start + length)

        return {
            "draw": draw,
            "recordsTotal": total_records,
            "recordsFiltered": filtered_records,
            "data": [
                self.output_instance(instance) for instance in query.all()
            ]
        }

    def output_instance(self, instance):
        returner = {
            key.name: self.get_value(key, instance) for key in self.columns
        }

        if self.data:
            returner["DT_RowData"] = {
                k: v(instance) for k, v in self.data.items()
            }

        return returner

    def get_value(self, key, instance):
        attr = key.model_name
        if "." in attr:
            tmp_list=attr.split(".")
            attr=tmp_list[-1]
            for sub in tmp_list[:-1]:
                instance = getattr(instance, sub)

        if key.filter is not None:
            r = key.filter(instance)
        else:
            r = getattr(instance, attr)

        return r() if inspect.isroutine(r) else r
<|endoftext|>