<commit_before>from __future__ import print_function
from peyotl import (assure_dir_exists,
                    get_logger,
                    write_as_json)
import codecs
import time
import os

_LOG = get_logger(__name__)

###################################################################################################
# Author: Stephen Smith

# Arguments:
#   download - T or F - whether or not to download the tar.gz file from NCBI
#   downloaddir - where to find (or put) the tar.gz and its contents
#   kill list file
#   destination dir - where taxonomy.tsv etc. are to be put

# JAR copied this file from data/ in the taxomachine repository
# to smasher/ in the opentree repository on 2013-04-25.
# Some subsequent modifications:
#  - remove "unclassified"
#  - add command line argument for directory in which to put ncbi
#  - change skipids from list to dictionary for speed
def file_mod_time_to_isotime(fp):
    sec_since_epoch = os.path.getmtime(fp)
    return epoch_seconds_to_isotime(sec_since_epoch)

def epoch_seconds_to_isotime(sec_since_epoch):
    tuple_time = time.gmtime(sec_since_epoch)
    return time.strftime("%Y-%m-%dT%H:%M:%S", tuple_time)

def add_or_append_to_dict(d, k, v):
    """If dict `d` has key `k`, then the new value of that key will be appended
    onto a list containing the previous values, otherwise d[k] = v.
    Creates a lightweight multimap, but not safe if v can be None or a list.
    returns True if the k now maps to >1 value"""
    ov = d.get(k)
    if ov is None:
        d[k] = v
        return False
    if isinstance(ov, list):
        ov.append(v)
    else:
        d[k] = [ov, v]
    return True

def parse_ncbi_names_file(names_fp):
    count = 0
    id_to_name = {}
    synonyms = {}
    names_to_ids = {}
    repeated_names = set()
    with codecs.open(names_fp, "r", encoding='utf-8') as namesf:
        for line in namesf:
            # if you do \t|\t then you don't get the name class right because it is "\t|"
            spls = line.split("\t|")
            node_id = int(spls[0])
            name = spls[1].strip()
            homonc = spls[2].strip()  # can get if it is a series here
            nm_c = spls[3].strip()  # scientific name, synonym, etc.
            if "<series>" in homonc:
                name = name + " series"
            if "subgroup <" in homonc:  # corrects some nested homonyms
                name = homonc.replace("<", "").replace(">", "")
            # nm_c can hold
            # scientific name   - the name used in OTT as primary.
            # synonym
            # equivalent name  - usually misspelling or spelling variant
            # misspelling
            # authority  - always extends scientific name
            # type material  - bacterial strain as type for prokaryotic species ??
            # common name
            # genbank common name
            # blast name   - 247 of them - a kind of common name
            # in-part (e.g. Bacteria in-part: Monera)
            # includes (what polarity?)
            if nm_c == "scientific name":
                id_to_name[node_id] = name
                if add_or_append_to_dict(names_to_ids, name, node_id):
                    repeated_names.add(name)
            elif nm_c != "in-part":
                synonyms.setdefault(node_id, []).append((node_id, name, nm_c))
            count += 1
            if count % 100000 == 0:
                _LOG.info('{} lines of names'.format(count))
    _LOG.info("number of lines in names file: {}".format(count))
    _LOG.info("number of distinct scientific names: {}".format(len(names_to_ids)))
    _LOG.info("number of synonyms: {}".format(len(synonyms)))
    return id_to_name, names_to_ids, synonyms, repeated_names

def deal_with_adj_taxa_with_same_names(id_to_parent,
                                       id_to_children,
                                       id_to_rank,
                                       id_to_name,
                                       names_to_ids,
                                       synonyms,
                                       repeated_names):
    """Here we look for cases in which a taxon and its parent have the same name.
    1. If the parent is a genus (this is common for subgenera), then we just alter the
        child's name to the form 'NAME CHILDS-RANK NAME'
    2. If the parent is not a genus, then we remove the child ID from the tree.
    """
    suppressed_ids, renamed_ids = set(), set()
    for name in repeated_names:
        ids_with_this_name = names_to_ids[name]
        assert isinstance(ids_with_this_name, list) and len(ids_with_this_name) > 1
        adj_same_named_ids = []
        for i in ids_with_this_name:
            par_id = id_to_parent.get(i)
            if par_id and par_id in ids_with_this_name:
                insert_loc = len(adj_same_named_ids)
                for ind, el in enumerate(adj_same_named_ids):
                    if el[1] == par_id:
                        insert_loc = ind
                        break
                adj_same_named_ids.insert(insert_loc, (par_id, i))
        for par_id, child_id in adj_same_named_ids:
            pr = id_to_rank.get(par_id)
            if pr and pr.lower() == 'genus':
                # Change the child's name
                cr = id_to_rank.get(child_id, '')
                nn = '{} {} {}'.format(name, cr, name)
                assert nn not in names_to_ids  # could happen, but ugh...
                names_to_ids[nn] = child_id
                id_to_name[child_id] = nn
                try:
                    ids_with_this_name.remove(child_id)
                except:
                    pass
                renamed_ids.add(child_id)
            else:
                # suppress the child
                suppressed_ids.add(child_id)
                c_list = id_to_children.get(child_id, [])
                pc_list = id_to_children.get(par_id)
                pc_list.remove(child_id)
                pc_list.extend(c_list)
                for gc in c_list:
                    id_to_parent[gc] = par_id
                for syn_el in synonyms.get(child_id, []):
                    synonyms.setdefault(par_id, []).append(syn_el)
    return suppressed_ids, renamed_ids

def parse_ncbi_nodes_file(nodes_fp):
    """Takes a filepath to an NCBI nodes.dmp and returns 3 dict mapping an ID to:
        - parent ID (can be None or an int)
        - children list (only for internals)
        - rank string (if available
    """
    count = 0
    to_par = {}  # key is the child id and the value is the parent
    to_children = {}  # key is the parent and value is the list of children
    to_rank = {}  # key is the node id and the value is the rank
    root_nodes = []
    with codecs.open(nodes_fp, "r", encoding='utf-8') as nodesf:
        for line in nodesf:
            spls = line.split("\t|\t")
            ns = spls[0].strip()
            node_id = int(ns)
            ps = spls[1].strip()
            # oddly enough, the root node is ID 1 and has parent ID 1 in nodes.dmp
            if ps and ps != ns:
                par_id = int(ps)
            else:
                par_id = None
                root_nodes.append(node_id)
            rank = spls[2].strip()
            to_par[node_id] = par_id
            if rank:
                to_rank[node_id] = rank
            to_children.setdefault(par_id, []).append(node_id)
            count += 1
            if count % 100000 == 0:
                _LOG.info('{} lines of nodes'.format(count))
    _LOG.info("number of lines in nodes file: {}".format(count))
    return to_par, to_children, to_rank, root_nodes

def write_ott_taxonomy(out_fp,
                       root_nodes,
                       id_to_par,
                       id_to_children,
                       id_to_rank,
                       id_to_name):
    with codecs.open(out_fp, 'w', encoding='utf-8') as out:
        out.write("uid\t|\tparent_uid\t|\tname\t|\trank\t|\t\n")
        # need to print id, parent id, and name
        for root_id in root_nodes:
            stack = [root_id]
            while stack:
                curr_id = stack.pop()
                name = id_to_name[curr_id]
                par_id = id_to_par[curr_id]
                spar_id = str(par_id)
                rank = id_to_rank.get(curr_id, '')
                children = id_to_children.get(curr_id)
                if children:
                    stack.extend(children)
                out.write('\t|\t'.join([str(curr_id), spar_id, name, rank]))
                out.write('\n')

def deal_with_ncbi_env_samples_names(id_to_par, id_to_name, names_to_ids):
    ess = 'environmental samples'
    es_ids = names_to_ids.setdefault(ess, [])
    renamed_ids = set(es_ids)
    for es_id in es_ids:
        par_id = id_to_par[es_id]
        name = '{} {}'.format(id_to_name[par_id], ess)
        id_to_name[es_id] = name
        add_or_append_to_dict(names_to_ids, name, es_id)
    del names_to_ids[ess]
    return renamed_ids

def normalize_ncbi(source, destination, url):
    nodes_fp = os.path.join(source, "nodes.dmp")
    names_fp = os.path.join(source, "names.dmp")
    about_obj = {"prefix": "ncbi",
                 "prefixDefinition": "http://www.ncbi.nlm.nih.gov/Taxonomy/Browser/wwwtax.cgi?id=",
                 "description": "NCBI Taxonomy",
                 "source": {"URL": url,
                            "date": file_mod_time_to_isotime(nodes_fp)
                            }
                 }
    about_fp = os.path.join(destination, 'about.json')
    id_to_par, id_to_children, id_to_rank, root_nodes = parse_ncbi_nodes_file(nodes_fp)
    assert len(root_nodes) == 1
    root_id = root_nodes[0]
    assert id_to_par[root_id] is None
    id_to_par[root_id] = ""
    id_to_name, name_to_ids, synonyms, repeated_names = parse_ncbi_names_file(names_fp)
    assert id_to_name[root_id] == "root"
    id_to_name[root_id] = "life"
    suppressed_ids, renamed_ids = deal_with_adj_taxa_with_same_names(id_to_par,
                                                                     id_to_children,
                                                                     id_to_rank,
                                                                     id_to_name,
                                                                     name_to_ids,
                                                                     synonyms,
                                                                     repeated_names)
    esr = deal_with_ncbi_env_samples_names(id_to_par, id_to_name, name_to_ids)
    assure_dir_exists(destination)
    write_ott_taxonomy(os.path.join(destination, 'taxonomy.tsv'),
                       root_nodes,
                       id_to_par,
                       id_to_children,
                       id_to_rank,
                       id_to_name)
    write_as_json(about_obj, about_fp, indent=2)

###################################################################################################
<commit_msg>basic ncbi normalization<commit_after>from __future__ import print_function

import codecs
import os
import time

from peyotl import (assure_dir_exists,
                    get_logger,
                    write_as_json)

_LOG = get_logger(__name__)

###################################################################################################
# Author: Stephen Smith

# Arguments:
#   download - T or F - whether or not to download the tar.gz file from NCBI
#   downloaddir - where to find (or put) the tar.gz and its contents
#   kill list file
#   destination dir - where taxonomy.tsv etc. are to be put

# JAR copied this file from data/ in the taxomachine repository
# to smasher/ in the opentree repository on 2013-04-25.
# Some subsequent modifications:
#  - remove "unclassified"
#  - add command line argument for directory in which to put ncbi
#  - change skipids from list to dictionary for speed

####################################################################################################
# utility
def file_mod_time_to_isotime(fp):
    sec_since_epoch = os.path.getmtime(fp)
    return epoch_seconds_to_isotime(sec_since_epoch)

def epoch_seconds_to_isotime(sec_since_epoch):
    tuple_time = time.gmtime(sec_since_epoch)
    return time.strftime("%Y-%m-%dT%H:%M:%S", tuple_time)

def add_or_append_to_dict(d, k, v):
    """If dict `d` has key `k`, then the new value of that key will be appended
    onto a list containing the previous values, otherwise d[k] = v.
    Creates a lightweight multimap, but not safe if v can be None or a list.
    returns True if the k now maps to >1 value"""
    ov = d.get(k)
    if ov is None:
        d[k] = v
        return False
    if isinstance(ov, list):
        ov.append(v)
    else:
        d[k] = [ov, v]
    return True

####################################################################################################
def parse_ncbi_names_file(names_fp, details_log):
    """Takes a filepath to an NCBI names.dmp file.
    Returns tuple
         0 id_to_name: node_id int -> str
         1 names_to_ids: str -> int or [int, ...]
         2 synonyms node_id -> [(name, type of synonym))
    """
    count = 0
    id_to_name = {}
    synonyms = {}
    names_to_ids = {}
    repeated_names = set()
    with codecs.open(names_fp, "r", encoding='utf-8') as namesf:
        for line in namesf:
            # if you do \t|\t then you don't get the name class right because it is "\t|"
            spls = line.split("\t|")
            node_id = int(spls[0])
            name = spls[1].strip()
            homonc = spls[2].strip()  # can get if it is a series here
            nm_c = spls[3].strip()  # scientific name, synonym, etc.
            if "<series>" in homonc:
                name = name + " series"
            if "subgroup <" in homonc:  # corrects some nested homonyms
                name = homonc.replace("<", "").replace(">", "")
            # nm_c can hold
            # scientific name   - the name used in OTT as primary.
            # synonym
            # equivalent name  - usually misspelling or spelling variant
            # misspelling
            # authority  - always extends scientific name
            # type material  - bacterial strain as type for prokaryotic species ??
            # common name
            # genbank common name
            # blast name   - 247 of them - a kind of common name
            # in-part (e.g. Bacteria in-part: Monera)
            # includes (what polarity?)
            if nm_c == "scientific name":
                id_to_name[node_id] = name
                if add_or_append_to_dict(names_to_ids, name, node_id):
                    repeated_names.add(name)
            elif nm_c != "in-part":
                synonyms.setdefault(node_id, []).append((name, nm_c))
            count += 1
            if count % 100000 == 0:
                _LOG.info('{} lines of names'.format(count))
    _LOG.info("number of lines in names file: {}".format(count))
    _LOG.info("number of distinct scientific names: {}".format(len(names_to_ids)))
    _LOG.info("number of IDs with synonyms: {}".format(len(synonyms)))
    details_log['num_distinct_names'] = len(names_to_ids)
    details_log['num_ids_with_synonyms'] = len(synonyms)
    return id_to_name, names_to_ids, synonyms, repeated_names

def parse_ncbi_nodes_file(nodes_fp, details_log):
    """Takes a filepath to an NCBI nodes.dmp and returns 3 dict mapping an ID to:
        - parent ID (can be None or an int)
        - children list (only for internals)
        - rank string (if available
    """
    count = 0
    to_par = {}  # key is the child id and the value is the parent
    to_children = {}  # key is the parent and value is the list of children
    to_rank = {}  # key is the node id and the value is the rank
    root_nodes = []
    with codecs.open(nodes_fp, "r", encoding='utf-8') as nodesf:
        for line in nodesf:
            spls = line.split("\t|\t")
            ns = spls[0].strip()
            node_id = int(ns)
            ps = spls[1].strip()
            # oddly enough, the root node is ID 1 and has parent ID 1 in nodes.dmp
            if ps and ps != ns:
                par_id = int(ps)
            else:
                par_id = None
                root_nodes.append(node_id)
            rank = spls[2].strip()
            to_par[node_id] = par_id
            if rank:
                to_rank[node_id] = rank
            to_children.setdefault(par_id, []).append(node_id)
            count += 1
            if count % 100000 == 0:
                _LOG.info('{} lines of nodes'.format(count))
    _LOG.info("number of lines in nodes file: {}".format(count))
    details_log['num_nodes'] = len(to_par)
    return to_par, to_children, to_rank, root_nodes

def deal_with_adj_taxa_with_same_names(id_to_parent,
                                       id_to_children,
                                       id_to_rank,
                                       id_to_name,
                                       names_to_ids,
                                       synonyms,
                                       repeated_names,
                                       details_log):
    """Here we look for cases in which a taxon and its parent have the same name.
    1. If the parent is a genus (this is common for subgenera), then we just alter the
        child's name to the form 'NAME CHILDS-RANK NAME'
    2. If the parent is not a genus, then we remove the child ID from the tree.
    """
    suppressed_ids, renamed_ids = {}, set()
    for name in repeated_names:
        ids_with_this_name = names_to_ids[name]
        assert isinstance(ids_with_this_name, list) and len(ids_with_this_name) > 1
        adj_same_named_ids = []
        for i in ids_with_this_name:
            par_id = id_to_parent.get(i)
            if par_id and par_id in ids_with_this_name:
                insert_loc = len(adj_same_named_ids)
                for ind, el in enumerate(adj_same_named_ids):
                    if el[1] == par_id:
                        insert_loc = ind
                        break
                adj_same_named_ids.insert(insert_loc, (par_id, i))
        for par_id, child_id in adj_same_named_ids:
            pr = id_to_rank.get(par_id)
            if pr and pr.lower() == 'genus':
                # Change the child's name
                cr = id_to_rank.get(child_id, '')
                nn = '{} {} {}'.format(name, cr, name)
                assert nn not in names_to_ids  # could happen, but ugh...
                names_to_ids[nn] = child_id
                id_to_name[child_id] = nn
                try:
                    ids_with_this_name.remove(child_id)
                except:
                    pass
                renamed_ids.add(child_id)
            else:
                # suppress the child
                suppressed_ids[child_id] = par_id
                c_list = id_to_children.get(child_id, [])
                pc_list = id_to_children.get(par_id)
                pc_list.remove(child_id)
                pc_list.extend(c_list)
                for gc in c_list:
                    id_to_parent[gc] = par_id
                for syn_el in synonyms.get(child_id, []):
                    synonyms.setdefault(par_id, []).append(syn_el)
    details_log['ids_suppressed_because_same_name_as_par'] = suppressed_ids
    ril = list(renamed_ids)
    ril.sort()
    details_log['names_decorated_because_same_name_as_par'] = ril
    return suppressed_ids, renamed_ids

def deal_with_ncbi_env_samples_names(id_to_par, id_to_name, names_to_ids, details_log):
    ess = 'environmental samples'
    es_ids = names_to_ids.setdefault(ess, [])
    renamed_ids = set(es_ids)
    for es_id in es_ids:
        par_id = id_to_par[es_id]
        name = '{} {}'.format(id_to_name[par_id], ess)
        id_to_name[es_id] = name
        add_or_append_to_dict(names_to_ids, name, es_id)
    del names_to_ids[ess]
    ril = list(renamed_ids)
    ril.sort()
    details_log['names_decorated_because_env_samp'] = ril
    return renamed_ids

####################################################################################################
def write_ott_taxonomy_tsv(out_fp,
                           root_nodes,
                           id_to_par,
                           id_to_children,
                           id_to_rank,
                           id_to_name,
                           has_syn_dict,
                           details_log):
    """If has_syn_dict is provided, then a list of the IDs that occur in that dict
    is returned in the order that the IDs were written to taxonomy file. This
    allows for the synonyms.tsv file to be written in a similar order, which makes browsing it
    easier.
    """
    if has_syn_dict is None:
        has_syn_dict = {}
    syn_id_order = []
    num_tips_written = 0
    num_internals_written = 0
    with codecs.open(out_fp, 'w', encoding='utf-8') as out:
        out.write("uid\t|\tparent_uid\t|\tname\t|\trank\t|\t\n")
        # need to print id, parent id, and name
        for root_id in root_nodes:
            stack = [root_id]
            while stack:
                curr_id = stack.pop()
                if curr_id in has_syn_dict:
                    syn_id_order.append(curr_id)
                name = id_to_name[curr_id]
                par_id = id_to_par[curr_id]
                spar_id = str(par_id)
                rank = id_to_rank.get(curr_id, '')
                children = id_to_children.get(curr_id)
                if children:
                    num_internals_written += 1
                    stack.extend(children)
                else:
                    num_tips_written += 1
                out.write('{}\n'.format('\t|\t'.join([str(curr_id), spar_id, name, rank])))
    details_log['num_tips_written'] = num_tips_written
    details_log['num_internals_written'] = num_internals_written
    return syn_id_order

def write_ott_synonyms_tsv(out_fp,
                           id_to_name_name_type_list,
                           id_order,
                           details_log):
    num_syn_written = 0
    with codecs.open(out_fp, 'w', encoding='utf-8') as out:
        out.write("uid\t|\tname\t|\ttype\t|\t\n")
        for nd_id in id_order:
            syn_list = id_to_name_name_type_list[nd_id]
            for name, name_type in syn_list:
                num_syn_written += 1
                out.write('{}\n'.format('\t|\t'.join([str(nd_id), name, name_type, ''])))
    details_log['num_synonyms_written'] = num_syn_written
    details_log['num_ids_with_synonyms_written'] = len(id_order)

def write_ncbi_details_json(fp, details_log):
    write_as_json(details_log, fp, indent=2)

def normalize_ncbi(source, destination, url):
    nodes_fp = os.path.join(source, "nodes.dmp")
    names_fp = os.path.join(source, "names.dmp")
    about_obj = {"prefix": "ncbi",
                 "prefixDefinition": "http://www.ncbi.nlm.nih.gov/Taxonomy/Browser/wwwtax.cgi?id=",
                 "description": "NCBI Taxonomy",
                 "source": {"URL": url,
                            "date": file_mod_time_to_isotime(nodes_fp)
                            }
                 }
    about_fp = os.path.join(destination, 'about.json')
    details_log = {}
    id_to_par, id_to_children, id_to_rank, root_nodes = parse_ncbi_nodes_file(nodes_fp,
                                                                              details_log)
    # Make sure there is only 1 root, and that its parent is an empty string
    assert len(root_nodes) == 1
    root_id = root_nodes[0]
    assert id_to_par[root_id] is None
    id_to_par[root_id] = ""
    id_to_name, name_to_ids, synonyms, repeated_names = parse_ncbi_names_file(names_fp,
                                                                              details_log)
    # Change the root's name from root to life
    assert id_to_name[root_id] == "root"
    id_to_name[root_id] = "life"
    # Clean up adjacent taxa with the same name
    # TODO: do we really want to suppress any? Seems dangerous... MTH Apr, 2017
    deal_with_adj_taxa_with_same_names(id_to_par,
                                       id_to_children,
                                       id_to_rank,
                                       id_to_name,
                                       name_to_ids,
                                       synonyms,
                                       repeated_names,
                                       details_log)
    # Decorate the names of environmental samples
    deal_with_ncbi_env_samples_names(id_to_par, id_to_name, name_to_ids, details_log)
    # Write out in OTT form
    assure_dir_exists(destination)
    syn_order = write_ott_taxonomy_tsv(os.path.join(destination, 'taxonomy.tsv'),
                                       root_nodes,
                                       id_to_par,
                                       id_to_children,
                                       id_to_rank,
                                       id_to_name,
                                       synonyms,
                                       details_log)
    write_ott_synonyms_tsv(os.path.join(destination, 'synonyms.tsv'),
                           synonyms,
                           syn_order,
                           details_log)
    write_ncbi_details_json(os.path.join(destination, 'details.json'),
                            details_log)
    write_as_json(about_obj, about_fp, indent=2)

###################################################################################################
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Author:  Dominik Gresch <greschd@gmx.ch>

import os
from collections.abc import Iterable
from functools import singledispatch

import click
import symmetry_representation as sr

from ._tb_model import Model

@click.group()
def cli():
    pass

def _output_option(**kwargs):
    return click.option(
        '--output', '-o',
        type=click.Path(dir_okay=False),
        **kwargs
    )

_input_option = click.option(
    '--input', '-i',
    type=click.Path(exists=True, dir_okay=False),
    default='model.hdf5',
    help='File containing the input model (in HDF5 format).'
)

def _read_input(input):
    click.echo("Reading initial model from file '{}' ...".format(input))
    return Model.from_hdf5_file(input)

def _write_output(model, output):
    click.echo("Writing output model to file '{}' ...".format(output))
    model.to_hdf5_file(output)
    click.echo("Done!")

@cli.command(short_help='Parse Wannier90 output files to an HDF5 file.')
@click.option(
    '--folder', '-f',
    type=click.Path(exists=True, file_okay=False),
    default='.',
    help='Directory containing the Wannier90 output files.'
)
@click.option(
    '--prefix', '-p',
    type=str,
    default='wannier',
    help='Common prefix of the Wannier90 output files.'
)
@_output_option(
    default='model.hdf5',
    help='Path of the output file.'
)
def parse(folder, prefix, output):
    """
    Parse Wannier90 output files and create an HDF5 file containing the tight-binding model.
    """
    click.echo("Parsing output files '{}*' ...".format(os.path.join(folder, prefix)))
    model = Model.from_wannier_folder(folder=folder, prefix=prefix, ignore_orbital_order=True)
    _write_output(model, output)

@cli.command(short_help='Create symmetrized tight-binding model.')
@_input_option
@_output_option(
    default='model_symmetrized.hdf5',
    help='Output file for the symmetrized model.'
)
@click.option(
    '--symmetries', '-s',
    type=click.Path(),
    default='symmetries.hdf5',
    help='File containing symmetry_representation.SymmetryGroup objects (in HDF5 form).'
)
@click.option(
    '--full-group/--no-full-group', '-f',
    default=None,
    help="""
    Full group: The full symmetry group is given in the symmetries.
    No full group: The symmetries only contain a generating subset of the full group. Overrides the option given in the symmetries file (if any).
    """
)
def symmetrize(input, output, symmetries, full_group):
    """
    Symmetrize tight-binding model with given symmetry group(s).
    """
    model = _read_input(input)
    click.echo("Reading symmetries from file '{}' ...".format(symmetries))
    sym = sr.io.load(symmetries)
    model_sym = _symmetrize(sym, model, full_group)
    _write_output(model_sym, output)

@singledispatch
def _symmetrize(sym, model, full_group):
    raise ValueError("Invalid type '{}' for _symmetrize".format(type(sym)))

@_symmetrize.register(Iterable)
def _(sym, model, full_group):
    for s in sym:
        model = _symmetrize(s, model, full_group)
    return model

@_symmetrize.register(sr.SymmetryGroup)
def _(sym, model, full_group):
    symmetries = sym.symmetries
    if full_group is None:
        full_group = sym.full_group
    click.echo("Symmetrizing model with {} symmetr{}, full_group={} ...".format(
        len(symmetries), 'y' if len(symmetries) == 1 else 'ies', full_group
    ))
    return model.symmetrize(symmetries=symmetries, full_group=full_group)

@_symmetrize.register(sr.SymmetryOperation)
def _(sym, model, full_group):
    sym_group = sr.SymmetryGroup(
        symmetries=[sym],
        full_group=full_group or False # catches 'None', does nothing for 'True' or 'False'
    )
    return _symmetrize(sym_group, model, full_group)

@cli.command(short_help="Slice specific orbitals from model.")
@_input_option
@_output_option(
    default='model_sliced.hdf5',
    help='Output file for the sliced model.'
)
@click.argument(
    'slice-idx',
    type=int,
    nargs=-1,
)
def slice(input, output, slice_idx):
    """
    Create a model containing only the orbitals given in the SLICE_IDX.
    """
    click.echo("Reading initial model from file '{}' ...".format(input))
    model = Model.from_hdf5_file(input)
    click.echo("Slicing model with indices {} ...".format(slice_idx))
    model_slice = model.slice_orbitals(slice_idx=slice_idx)
    _write_output(model_slice, output)
<commit_msg>Use _read_input in slice CL implementation<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Author:  Dominik Gresch <greschd@gmx.ch>

import os
from collections.abc import Iterable
from functools import singledispatch

import click
import symmetry_representation as sr

from ._tb_model import Model

@click.group()
def cli():
    pass

def _output_option(**kwargs):
    return click.option(
        '--output', '-o',
        type=click.Path(dir_okay=False),
        **kwargs
    )

_input_option = click.option(
    '--input', '-i',
    type=click.Path(exists=True, dir_okay=False),
    default='model.hdf5',
    help='File containing the input model (in HDF5 format).'
)

def _read_input(input):
    click.echo("Reading initial model from file '{}' ...".format(input))
    return Model.from_hdf5_file(input)

def _write_output(model, output):
    click.echo("Writing output model to file '{}' ...".format(output))
    model.to_hdf5_file(output)
    click.echo("Done!")

@cli.command(short_help='Parse Wannier90 output files to an HDF5 file.')
@click.option(
    '--folder', '-f',
    type=click.Path(exists=True, file_okay=False),
    default='.',
    help='Directory containing the Wannier90 output files.'
)
@click.option(
    '--prefix', '-p',
    type=str,
    default='wannier',
    help='Common prefix of the Wannier90 output files.'
)
@_output_option(
    default='model.hdf5',
    help='Path of the output file.'
)
def parse(folder, prefix, output):
    """
    Parse Wannier90 output files and create an HDF5 file containing the tight-binding model.
    """
    click.echo("Parsing output files '{}*' ...".format(os.path.join(folder, prefix)))
    model = Model.from_wannier_folder(folder=folder, prefix=prefix, ignore_orbital_order=True)
    _write_output(model, output)

@cli.command(short_help='Create symmetrized tight-binding model.')
@_input_option
@_output_option(
    default='model_symmetrized.hdf5',
    help='Output file for the symmetrized model.'
)
@click.option(
    '--symmetries', '-s',
    type=click.Path(),
    default='symmetries.hdf5',
    help='File containing symmetry_representation.SymmetryGroup objects (in HDF5 form).'
)
@click.option(
    '--full-group/--no-full-group', '-f',
    default=None,
    help="""
    Full group: The full symmetry group is given in the symmetries.
    No full group: The symmetries only contain a generating subset of the full group. Overrides the option given in the symmetries file (if any).
    """
)
def symmetrize(input, output, symmetries, full_group):
    """
    Symmetrize tight-binding model with given symmetry group(s).
    """
    model = _read_input(input)
    click.echo("Reading symmetries from file '{}' ...".format(symmetries))
    sym = sr.io.load(symmetries)
    model_sym = _symmetrize(sym, model, full_group)
    _write_output(model_sym, output)

@singledispatch
def _symmetrize(sym, model, full_group):
    raise ValueError("Invalid type '{}' for _symmetrize".format(type(sym)))

@_symmetrize.register(Iterable)
def _(sym, model, full_group):
    for s in sym:
        model = _symmetrize(s, model, full_group)
    return model

@_symmetrize.register(sr.SymmetryGroup)
def _(sym, model, full_group):
    symmetries = sym.symmetries
    if full_group is None:
        full_group = sym.full_group
    click.echo("Symmetrizing model with {} symmetr{}, full_group={} ...".format(
        len(symmetries), 'y' if len(symmetries) == 1 else 'ies', full_group
    ))
    return model.symmetrize(symmetries=symmetries, full_group=full_group)

@_symmetrize.register(sr.SymmetryOperation)
def _(sym, model, full_group):
    sym_group = sr.SymmetryGroup(
        symmetries=[sym],
        full_group=full_group or False # catches 'None', does nothing for 'True' or 'False'
    )
    return _symmetrize(sym_group, model, full_group)

@cli.command(short_help="Slice specific orbitals from model.")
@_input_option
@_output_option(
    default='model_sliced.hdf5',
    help='Output file for the sliced model.'
)
@click.argument(
    'slice-idx',
    type=int,
    nargs=-1,
)
def slice(input, output, slice_idx):
    """
    Create a model containing only the orbitals given in the SLICE_IDX.
    """
    model = _read_input(input)
    click.echo("Slicing model with indices {} ...".format(slice_idx))
    model_slice = model.slice_orbitals(slice_idx=slice_idx)
    _write_output(model_slice, output)
<|endoftext|>


<commit_before>############################################################################
## Tool name: BetterBusBuffers - Count Trips at Stops by Route and Direction
## Created by: David Wasserman, https://github.com/d-wasserman and Melinda Morang, Esri
## This tool was developed as part of Transit R&D Efforts from Fehr & Peers.
## Fehr & Peers contributes this tool to the BBB Toolset to further more
## informed planning.
## Last updated: 8 October 2020
############################################################################
''' BetterBusBuffers - Count Trips at Stops by Route and Direction

BetterBusBuffers provides a quantitative measure of access to public transit
in your city by counting the transit trip frequency at various locations.

The Count Trips at Stops by Route and Direction outputs a feature class where
every GTFS stop is duplicated for every route-direction combination that uses
that stop during the analysis time windows. Each point will represent a unique
combination of stop id, route id, and direction id, and the frequency statistics
that relate to each of them for the analyzed time window.
'''
################################################################################
'''Copyright 2020 Esri
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
       http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.'''

################################################################################
"""Copyright 2020 Fehr & Peers

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
################################################################################
"""Copyright 2020 David Wasserman

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
################################################################################
import arcpy
import BBB_SharedFunctions
import sqlite3, os, datetime

def GenerateTimePeriodExtent(start_time_stamp_list, end_time_stamp_list):
    """Takes a list of start and end time stamps and returns the minimum valued start time and the maximum valued
    end time to determine the global temporal extent."""
    temporal_extent = None
    start_seconds_list = []
    end_seconds_list = []
    for start_time, end_time in zip(start_time_stamp_list, end_time_stamp_list):
        # Convert to seconds from midnight
        start_sec = BBB_SharedFunctions.parse_time(start_time + ":00")
        # Convert to seconds from midnight
        end_sec = BBB_SharedFunctions.parse_time(end_time + ":00")
        start_seconds_list.append(start_sec)
        end_seconds_list.append(end_sec)
    min_start_time, max_end_time = min(start_seconds_list), max(end_seconds_list)
    return min_start_time, max_end_time

def GenerateTimePeriodList(start_time_stamp_list, end_time_stamp_list, alias_list, start_default="00:00",
                           end_default="23:59"):
    """This function will take a list of string time stamps of start and end times, and return a nested list of the structure.
    [(Time_period_iD, start_seconds, end_seconds, time_window),...]
    :param - start_time_stamp_list - list of start times in HH:MM
    :param - end_time_stamp_list - list of end times in HH:MM
    :param - alias_list- list of names to name time period summaries
    :param - start_default - default start time
    :param - end_default -default end time"""
    time_period_list = []
    cached_hours = (None, None)
    for start_time, end_time, alias in zip(start_time_stamp_list, end_time_stamp_list, alias_list):
        # Declare Defaults
        if start_time == "":
            start_time = start_default
        # Convert to seconds
        start_sec = BBB_SharedFunctions.parse_time(start_time + ":00")
        # Upper end of time window (HH:MM in 24-hour time)
        # Default end time is 11:59pm if they leave it blank.
        if end_time == "":
            end_time = end_default
        # Generate Unique Time Period ID
        start_time_id = str(start_time).split(":")[0]
        end_time_id = str(end_time).split(":")[0]
        if start_time_id == cached_hours[0] and end_time_id == cached_hours[1]:
            start_time_id = str(start_time).replace(":", "")
            end_time_id = str(end_time).replace(":", "")
        finalID = alias
        cached_hours = (start_time_id, end_time_id)
        # Determine Temporal Values
        # Convert to seconds
        end_sec = BBB_SharedFunctions.parse_time(end_time + ":00")
        # Length of time window in hours - used to calculate num trips / hour
        # TimeWindowLength = (end_sec - start_sec) / 3600
        TimeWindowLength = (end_sec - start_sec) / 3600
        time_period_list.append((finalID, start_sec, end_sec, TimeWindowLength))
    return time_period_list

def runTool(output_stop_file, SQLDbase, time_window_value_table, snap_to_nearest_5_minutes):
    def RetrieveFrequencyStatsForStop(stop_id, stoptimedict, start_sec, end_sec):
        '''For a given stop, query the dictionary
        and return the NumTrips, NumTripsPerHr, MaxWaitTime, and AvgHeadway given a
        specific route_id and direction. If snap to nearest five minutes is true, then
        this function will return headways snapped to the closest 5 minute interval.'''
        # Make a list of stop_times
        StopTimesAtThisPoint = []
        try:
            for trip in stoptimedict[stop_id]:
                StopTimesAtThisPoint.append(trip[1])
        except KeyError:
            pass
        StopTimesAtThisPoint.sort()

        # Calculate the number of trips
        NumTrips = len(StopTimesAtThisPoint)
        NumTripsPerHr = round(float(NumTrips) / ((end_sec - start_sec) / 3600), 2)
        # Get the max wait time and the average headway
        MaxWaitTime = BBB_SharedFunctions.CalculateMaxWaitTime(StopTimesAtThisPoint, start_sec, end_sec)
        if snap_to_nearest_5_minutes:
            round_to = 5
        else:
            round_to = None
        AvgHeadway = BBB_SharedFunctions.CalculateAvgHeadway(StopTimesAtThisPoint, round_to)
        return NumTrips, NumTripsPerHr, MaxWaitTime, AvgHeadway

    # ----- Get input parameters and set things up. -----
    # Check software version and fail out quickly if it's not sufficient.
    BBB_SharedFunctions.CheckArcVersion(min_version_pro="1.2")

    arcpy.AddMessage("Reading data...")

    # Connect to SQL database of preprocessed GTFS from Step 1
    conn = BBB_SharedFunctions.conn = sqlite3.connect(SQLDbase)
    c = BBB_SharedFunctions.c = conn.cursor()

    # Store frequencies if relevant
    frequencies_dict = BBB_SharedFunctions.MakeFrequenciesDict()

    # Get unique route_id/direction_id pairs and calculate the trips used in each
    # Some GTFS datasets use the same route_id to identify trips traveling in
    # either direction along a route. Others identify it as a different route.
    # We will consider each direction separately if there is more than one.
    trip_route_dict = {}  # {(route_id, direction_id): [(trip_id, service_id),..]}
    triproutefetch = '''SELECT DISTINCT route_id,direction_id FROM trips;'''
    c.execute(triproutefetch)
    for rtpair in c.fetchall():
        key = tuple(rtpair)
        route_id = rtpair[0]
        direction_id = rtpair[1]
        # Get list of trips
        # Ignore direction if this route doesn't have a direction
        if direction_id is not None and str(direction_id).strip():
            triproutefetch = '''
                    SELECT trip_id, service_id FROM trips
                    WHERE route_id = '{0}' AND direction_id = {1};'''.format(route_id, direction_id)
        else:
            triproutefetch = '''
                    SELECT trip_id, service_id FROM trips
                    WHERE route_id = '{0}';'''.format(route_id)
        c.execute(triproutefetch)
        triproutelist = c.fetchall()
        trip_route_dict[key] = triproutelist

    # ----- For each time window, calculate the stop frequency -----
    final_stop_freq_dict = {}  # {(stop_id, route_id, direction_id): {prefix: (NumTrips, NumTripsPerHour, MaxWaitTimeSec, AvgHeadwayMin)}}
    # The time_window_value_table will be a list of nested lists of strings like:
    # [[Weekday name or YYYYMMDD date, HH: MM, HH: MM, Departures / Arrivals, Prefix], [], ...]
    for time_window in time_window_value_table:
        # Prefix/identifier associated with this time window
        prefix = time_window[4]
        arcpy.AddMessage("Calculating statistics for time window %s..." % prefix)
        # Clean up date and determine whether it's a date or a weekday
        Specific, day = BBB_SharedFunctions.CheckSpecificDate(time_window[0])
        # Convert times to seconds
        start_time = time_window[1]
        end_time = time_window[2]
        if not start_time:
            start_time = "00:00"
        if not end_time:
            end_time = "23:59"
        start_sec, end_sec = BBB_SharedFunctions.ConvertTimeWindowToSeconds(start_time, end_time)
        # Clean up arrival/departure time choice
        DepOrArr = BBB_SharedFunctions.CleanUpDepOrArr(time_window[3])

        # Get the trips running in this time window for each route/direction pair
        # Get the service_ids serving the correct days
        serviceidlist, serviceidlist_yest, serviceidlist_tom = \
            BBB_SharedFunctions.GetServiceIDListsAndNonOverlaps(day, start_sec, end_sec, DepOrArr, Specific)

        # Retrieve the stop_times for the time window broken out by route/direction
        stoproutedir_dict = {}  # {(stop_id, route_id, direction_id): [NumTrips, NumTripsPerHour, MaxWaitTimeSec, AvgHeadwayMin]}
        for rtdirpair in trip_route_dict:
            # Get trips running with these service_ids
            trip_serv_list = trip_route_dict[rtdirpair]
            triplist = []
            for tripserv in trip_serv_list:
                # Only keep trips running on the correct day
                if tripserv[1] in serviceidlist or tripserv[1] in serviceidlist_tom or \
                    tripserv[1] in serviceidlist_yest:
                    triplist.append(tripserv[0])

            # Get the stop_times that occur during this time window for these trips
            try:
                stoptimedict = BBB_SharedFunctions.GetStopTimesForStopsInTimeWindow(
                    start_sec, end_sec, DepOrArr, triplist, "today", frequencies_dict)
            except KeyError:  # No trips
                pass
            try:
                stoptimedict_yest = BBB_SharedFunctions.GetStopTimesForStopsInTimeWindow(
                    start_sec, end_sec, DepOrArr, triplist, "yesterday", frequencies_dict)
            except KeyError:  # No trips
                pass
            try:
                stoptimedict_tom = BBB_SharedFunctions.GetStopTimesForStopsInTimeWindow(
                    start_sec, end_sec, DepOrArr, triplist, "tomorrow", frequencies_dict)
            except KeyError:  # No trips
                pass

            # Combine the three dictionaries into one master
            for stop in stoptimedict_yest:
                stoptimedict[stop] = stoptimedict.setdefault(stop, []) + stoptimedict_yest[stop]
            for stop in stoptimedict_tom:
                stoptimedict[stop] = stoptimedict.setdefault(stop, []) + stoptimedict_tom[stop]

            for stop in stoptimedict.keys():
                # Get Stop-Route-Dir Frequencies by time period
                vals = RetrieveFrequencyStatsForStop(stop, stoptimedict, start_sec, end_sec)
                key = (stop, rtdirpair[0], rtdirpair[1],)
                if key not in final_stop_freq_dict:
                    final_stop_freq_dict[key] = {prefix: vals}
                else:
                    final_stop_freq_dict[key][prefix] = vals

    # ----- Write the stops and stats to the output feature class -----
    arcpy.AddMessage("Writing outputs...")
    # Make the basic feature class for stops with correct gtfs fields
    with arcpy.EnvManager(overwriteOutput=True):
        output_coords = BBB_SharedFunctions.CreateStopsFeatureClass(output_stop_file)

    # Add fields specific to this tool's outputs
    arcpy.management.AddField(output_stop_file, 'route_id', "TEXT")
    arcpy.management.AddField(output_stop_file, 'direction_id', "SHORT")
    # Create fields for stats for each time window using prefix
    base_field_names = ['_NumTrips', '_NumTripsPerHr', '_MaxWaitTime', '_AvgHeadway']
    new_fields = []
    for time_window in time_window_value_table:
        for base_field in base_field_names:
            new_field = time_window[4] + base_field
            new_fields.append(new_field)
            arcpy.management.AddField(output_stop_file, new_field, "DOUBLE")

    # Get the stop info from the GTFS SQL file
    StopTable = BBB_SharedFunctions.GetStopsData()
    stop_dict = {stop[0]: stop for stop in StopTable}

    # Make a dictionary to track whether we have inserted all stops at least once into the output
    used_stops = {stop[0]: False for stop in StopTable}
    # Store stop geometries in dictionary so they can be inserted multiple times without recalculating
    stop_geoms = {stop[0]: BBB_SharedFunctions.MakeStopGeometry(stop[4], stop[5], output_coords) for stop in StopTable}

    # Add the stops with stats to the feature class
    fields = [
        "SHAPE@", "stop_id", "stop_code", "stop_name", "stop_desc", "zone_id", "stop_url", "location_type",
        "parent_station", "route_id", "direction_id"
    ] + new_fields
    with arcpy.da.InsertCursor(output_stop_file, fields) as cur3:
        # Iterate over all unique stop, route_id, direction_id groups and insert values
        for key in sorted(final_stop_freq_dict.keys()):
            stop_id = key[0]
            used_stops[stop_id] = True
            route_id = key[1]
            direction_id = key[2]
            stop_data = stop_dict[stop_id]
            # Schema of StopTable
            ##   0 - stop_id
            ##   1 - stop_code
            ##   2 - stop_name
            ##   3 - stop_desc
            ##   4 - stop_lat
            ##   5 - stop_lon
            ##   6 - zone_id
            ##   7 - stop_url
            ##   8 - location_type
            ##   9 - parent_station
            row = [
                stop_geoms[stop_id],  # Geometry
                stop_data[0], stop_data[1], stop_data[2], stop_data[3], stop_data[6], stop_data[7], stop_data[8], stop_data[9],  # GTFS data
                route_id, direction_id  # route and direction IDs
            ]
            # Populate stats fields for each prefix
            for time_window in time_window_value_table:
                prefix = time_window[4]
                try:
                    vals = final_stop_freq_dict[key][prefix]
                except KeyError:
                    # This stop/route/direction group had no service for this time window
                    vals = [0, 0, None, None]
                row += vals

            # Insert the row
            cur3.insertRow(row)

        # Insert row for any remaining stops that were not used at all
        for stop_id in used_stops:
            if used_stops[stop_id]:
                # This one was already inserted
                continue
            stop_data = stop_dict[stop_id]
            row = [
                stop_geoms[stop_id],  # Geometry
                stop_data[0], stop_data[1], stop_data[2], stop_data[3], stop_data[6], stop_data[7], stop_data[8], stop_data[9],  # GTFS data
                None, None  # route and direction IDs - None because not used
            ]
            # Populate stats fields for each prefix
            for time_window in time_window_value_table:
                row += [0, 0, None, None]
            # Insert the row
            cur3.insertRow(row)

    # Close Connection
    conn.close()
    arcpy.AddMessage("Finished!")
    arcpy.AddMessage("Calculated trip counts, frequency, max wait time, and \
headway were written to an output stops file by route-direction pairs.")
<commit_msg>Remove no-longer-used functions<commit_after>############################################################################
## Tool name: BetterBusBuffers - Count Trips at Stops by Route and Direction
## Created by: David Wasserman, https://github.com/d-wasserman and Melinda Morang, Esri
## This tool was developed as part of Transit R&D Efforts from Fehr & Peers.
## Fehr & Peers contributes this tool to the BBB Toolset to further more
## informed planning.
## Last updated: 8 October 2020
############################################################################
''' BetterBusBuffers - Count Trips at Stops by Route and Direction

BetterBusBuffers provides a quantitative measure of access to public transit
in your city by counting the transit trip frequency at various locations.

The Count Trips at Stops by Route and Direction outputs a feature class where
every GTFS stop is duplicated for every route-direction combination that uses
that stop during the analysis time windows. Each point will represent a unique
combination of stop id, route id, and direction id, and the frequency statistics
that relate to each of them for the analyzed time window.
'''
################################################################################
'''Copyright 2020 Esri
   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at
       http://www.apache.org/licenses/LICENSE-2.0
   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.'''

################################################################################
"""Copyright 2020 Fehr & Peers

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
################################################################################
"""Copyright 2020 David Wasserman

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""
################################################################################
import arcpy
import BBB_SharedFunctions
import sqlite3, os, datetime

def runTool(output_stop_file, SQLDbase, time_window_value_table, snap_to_nearest_5_minutes):
    def RetrieveFrequencyStatsForStop(stop_id, stoptimedict, start_sec, end_sec):
        '''For a given stop, query the dictionary
        and return the NumTrips, NumTripsPerHr, MaxWaitTime, and AvgHeadway given a
        specific route_id and direction. If snap to nearest five minutes is true, then
        this function will return headways snapped to the closest 5 minute interval.'''
        # Make a list of stop_times
        StopTimesAtThisPoint = []
        try:
            for trip in stoptimedict[stop_id]:
                StopTimesAtThisPoint.append(trip[1])
        except KeyError:
            pass
        StopTimesAtThisPoint.sort()

        # Calculate the number of trips
        NumTrips = len(StopTimesAtThisPoint)
        NumTripsPerHr = round(float(NumTrips) / ((end_sec - start_sec) / 3600), 2)
        # Get the max wait time and the average headway
        MaxWaitTime = BBB_SharedFunctions.CalculateMaxWaitTime(StopTimesAtThisPoint, start_sec, end_sec)
        if snap_to_nearest_5_minutes:
            round_to = 5
        else:
            round_to = None
        AvgHeadway = BBB_SharedFunctions.CalculateAvgHeadway(StopTimesAtThisPoint, round_to)
        return NumTrips, NumTripsPerHr, MaxWaitTime, AvgHeadway

    # ----- Get input parameters and set things up. -----
    # Check software version and fail out quickly if it's not sufficient.
    BBB_SharedFunctions.CheckArcVersion(min_version_pro="1.2")

    arcpy.AddMessage("Reading data...")

    # Connect to SQL database of preprocessed GTFS from Step 1
    conn = BBB_SharedFunctions.conn = sqlite3.connect(SQLDbase)
    c = BBB_SharedFunctions.c = conn.cursor()

    # Store frequencies if relevant
    frequencies_dict = BBB_SharedFunctions.MakeFrequenciesDict()

    # Get unique route_id/direction_id pairs and calculate the trips used in each
    # Some GTFS datasets use the same route_id to identify trips traveling in
    # either direction along a route. Others identify it as a different route.
    # We will consider each direction separately if there is more than one.
    trip_route_dict = {}  # {(route_id, direction_id): [(trip_id, service_id),..]}
    triproutefetch = '''SELECT DISTINCT route_id,direction_id FROM trips;'''
    c.execute(triproutefetch)
    for rtpair in c.fetchall():
        key = tuple(rtpair)
        route_id = rtpair[0]
        direction_id = rtpair[1]
        # Get list of trips
        # Ignore direction if this route doesn't have a direction
        if direction_id is not None and str(direction_id).strip():
            triproutefetch = '''
                    SELECT trip_id, service_id FROM trips
                    WHERE route_id = '{0}' AND direction_id = {1};'''.format(route_id, direction_id)
        else:
            triproutefetch = '''
                    SELECT trip_id, service_id FROM trips
                    WHERE route_id = '{0}';'''.format(route_id)
        c.execute(triproutefetch)
        triproutelist = c.fetchall()
        trip_route_dict[key] = triproutelist

    # ----- For each time window, calculate the stop frequency -----
    final_stop_freq_dict = {}  # {(stop_id, route_id, direction_id): {prefix: (NumTrips, NumTripsPerHour, MaxWaitTimeSec, AvgHeadwayMin)}}
    # The time_window_value_table will be a list of nested lists of strings like:
    # [[Weekday name or YYYYMMDD date, HH: MM, HH: MM, Departures / Arrivals, Prefix], [], ...]
    for time_window in time_window_value_table:
        # Prefix/identifier associated with this time window
        prefix = time_window[4]
        arcpy.AddMessage("Calculating statistics for time window %s..." % prefix)
        # Clean up date and determine whether it's a date or a weekday
        Specific, day = BBB_SharedFunctions.CheckSpecificDate(time_window[0])
        # Convert times to seconds
        start_time = time_window[1]
        end_time = time_window[2]
        if not start_time:
            start_time = "00:00"
        if not end_time:
            end_time = "23:59"
        start_sec, end_sec = BBB_SharedFunctions.ConvertTimeWindowToSeconds(start_time, end_time)
        # Clean up arrival/departure time choice
        DepOrArr = BBB_SharedFunctions.CleanUpDepOrArr(time_window[3])

        # Get the trips running in this time window for each route/direction pair
        # Get the service_ids serving the correct days
        serviceidlist, serviceidlist_yest, serviceidlist_tom = \
            BBB_SharedFunctions.GetServiceIDListsAndNonOverlaps(day, start_sec, end_sec, DepOrArr, Specific)

        # Retrieve the stop_times for the time window broken out by route/direction
        stoproutedir_dict = {}  # {(stop_id, route_id, direction_id): [NumTrips, NumTripsPerHour, MaxWaitTimeSec, AvgHeadwayMin]}
        for rtdirpair in trip_route_dict:
            # Get trips running with these service_ids
            trip_serv_list = trip_route_dict[rtdirpair]
            triplist = []
            for tripserv in trip_serv_list:
                # Only keep trips running on the correct day
                if tripserv[1] in serviceidlist or tripserv[1] in serviceidlist_tom or \
                    tripserv[1] in serviceidlist_yest:
                    triplist.append(tripserv[0])

            # Get the stop_times that occur during this time window for these trips
            try:
                stoptimedict = BBB_SharedFunctions.GetStopTimesForStopsInTimeWindow(
                    start_sec, end_sec, DepOrArr, triplist, "today", frequencies_dict)
            except KeyError:  # No trips
                pass
            try:
                stoptimedict_yest = BBB_SharedFunctions.GetStopTimesForStopsInTimeWindow(
                    start_sec, end_sec, DepOrArr, triplist, "yesterday", frequencies_dict)
            except KeyError:  # No trips
                pass
            try:
                stoptimedict_tom = BBB_SharedFunctions.GetStopTimesForStopsInTimeWindow(
                    start_sec, end_sec, DepOrArr, triplist, "tomorrow", frequencies_dict)
            except KeyError:  # No trips
                pass

            # Combine the three dictionaries into one master
            for stop in stoptimedict_yest:
                stoptimedict[stop] = stoptimedict.setdefault(stop, []) + stoptimedict_yest[stop]
            for stop in stoptimedict_tom:
                stoptimedict[stop] = stoptimedict.setdefault(stop, []) + stoptimedict_tom[stop]

            for stop in stoptimedict.keys():
                # Get Stop-Route-Dir Frequencies by time period
                vals = RetrieveFrequencyStatsForStop(stop, stoptimedict, start_sec, end_sec)
                key = (stop, rtdirpair[0], rtdirpair[1],)
                if key not in final_stop_freq_dict:
                    final_stop_freq_dict[key] = {prefix: vals}
                else:
                    final_stop_freq_dict[key][prefix] = vals

    # ----- Write the stops and stats to the output feature class -----
    arcpy.AddMessage("Writing outputs...")
    # Make the basic feature class for stops with correct gtfs fields
    with arcpy.EnvManager(overwriteOutput=True):
        output_coords = BBB_SharedFunctions.CreateStopsFeatureClass(output_stop_file)

    # Add fields specific to this tool's outputs
    arcpy.management.AddField(output_stop_file, 'route_id', "TEXT")
    arcpy.management.AddField(output_stop_file, 'direction_id', "SHORT")
    # Create fields for stats for each time window using prefix
    base_field_names = ['_NumTrips', '_NumTripsPerHr', '_MaxWaitTime', '_AvgHeadway']
    new_fields = []
    for time_window in time_window_value_table:
        for base_field in base_field_names:
            new_field = time_window[4] + base_field
            new_fields.append(new_field)
            arcpy.management.AddField(output_stop_file, new_field, "DOUBLE")

    # Get the stop info from the GTFS SQL file
    StopTable = BBB_SharedFunctions.GetStopsData()
    stop_dict = {stop[0]: stop for stop in StopTable}

    # Make a dictionary to track whether we have inserted all stops at least once into the output
    used_stops = {stop[0]: False for stop in StopTable}
    # Store stop geometries in dictionary so they can be inserted multiple times without recalculating
    stop_geoms = {stop[0]: BBB_SharedFunctions.MakeStopGeometry(stop[4], stop[5], output_coords) for stop in StopTable}

    # Add the stops with stats to the feature class
    fields = [
        "SHAPE@", "stop_id", "stop_code", "stop_name", "stop_desc", "zone_id", "stop_url", "location_type",
        "parent_station", "route_id", "direction_id"
    ] + new_fields
    with arcpy.da.InsertCursor(output_stop_file, fields) as cur3:
        # Iterate over all unique stop, route_id, direction_id groups and insert values
        for key in sorted(final_stop_freq_dict.keys()):
            stop_id = key[0]
            used_stops[stop_id] = True
            route_id = key[1]
            direction_id = key[2]
            stop_data = stop_dict[stop_id]
            # Schema of StopTable
            ##   0 - stop_id
            ##   1 - stop_code
            ##   2 - stop_name
            ##   3 - stop_desc
            ##   4 - stop_lat
            ##   5 - stop_lon
            ##   6 - zone_id
            ##   7 - stop_url
            ##   8 - location_type
            ##   9 - parent_station
            row = [
                stop_geoms[stop_id],  # Geometry
                stop_data[0], stop_data[1], stop_data[2], stop_data[3], stop_data[6], stop_data[7], stop_data[8], stop_data[9],  # GTFS data
                route_id, direction_id  # route and direction IDs
            ]
            # Populate stats fields for each prefix
            for time_window in time_window_value_table:
                prefix = time_window[4]
                try:
                    vals = final_stop_freq_dict[key][prefix]
                except KeyError:
                    # This stop/route/direction group had no service for this time window
                    vals = [0, 0, None, None]
                row += vals

            # Insert the row
            cur3.insertRow(row)

        # Insert row for any remaining stops that were not used at all
        for stop_id in used_stops:
            if used_stops[stop_id]:
                # This one was already inserted
                continue
            stop_data = stop_dict[stop_id]
            row = [
                stop_geoms[stop_id],  # Geometry
                stop_data[0], stop_data[1], stop_data[2], stop_data[3], stop_data[6], stop_data[7], stop_data[8], stop_data[9],  # GTFS data
                None, None  # route and direction IDs - None because not used
            ]
            # Populate stats fields for each prefix
            for time_window in time_window_value_table:
                row += [0, 0, None, None]
            # Insert the row
            cur3.insertRow(row)

    # Close Connection
    conn.close()
    arcpy.AddMessage("Finished!")
    arcpy.AddMessage("Calculated trip counts, frequency, max wait time, and \
headway were written to an output stops file by route-direction pairs.")
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import sqlite3
import os

from pkg_resources import resource_string

from csv_parser import parse_csv
import reports

DB_NAME = 'tellfl.db'

INSERT_PAYMENT = 'INSERT INTO payments(user, amount, time) VALUES(?, ?, ?)'
INSERT_JOURNEY = 'INSERT INTO journeys(user, station_from, station_to, time_in, time_out, cost) VALUES(?, ?, ?, ?, ?, ?)' # NOQA

def find_user(cursor, username):
    q = 'SELECT id FROM users WHERE username=?'
    cursor.execute(q, (username, ))
    try:
        return int(cursor.fetchone()[0])
    except:
        raise Exception('No user with that username')

def add_to_db(cursor, uid, csv_path):
    for record in parse_csv(csv_path):
        if record['type'] == 'payment':
            cursor.execute(
                INSERT_PAYMENT,
                (uid, record['amount'], record['time'])
            )
        elif record['type'] == 'journey':
            cursor.execute(
                INSERT_JOURNEY,
                (uid, record['station_from'], record['station_to'],
                 record['time_in'], record['time_out'], record['cost'])
            )

class TellfLException(Exception):
    pass

class DB:

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

    def __enter__(self):
        self.conn = sqlite3.connect(os.path.join(self.location, DB_NAME))
        return self.conn.cursor()

    def __exit__(self, type, value, traceback):
        if traceback is None:
            self.conn.commit()
        self.conn.close()

def report(username, location, as_html):
    with DB(location) as cursor:
        uid = find_user(cursor, username)
        r = reports.create(cursor, uid)
        if as_html:
            print reports.asHtml(r)
        else:
            print reports.asText(r)

def add(username, csv_path, location):
    with DB(location) as cursor:
        uid = find_user(cursor, username)
        add_to_db(cursor, uid, csv_path)

def register(username, location):
    with DB(location) as cursor:
        try:
            q = 'INSERT INTO users(username) VALUES(?)'
            cursor.execute(q, (username, ))
        except:
            raise TellfLException('A user with that username already exists')

def install(location):
    os.makedirs(location)
    with DB(location) as cursor:
        sql = resource_string(__name__, 'assets/schema.sql')
        cursor.executescript(sql)
<commit_msg>formatting<commit_after># -*- coding: utf-8 -*-
import sqlite3
import os

from pkg_resources import resource_string

from csv_parser import parse_csv
import reports

DB_NAME = 'tellfl.db'

INSERT_PAYMENT = ('INSERT INTO '
                  '    payments(user, amount, time) '
                  'VALUES(?, ?, ?)')
INSERT_JOURNEY = ('INSERT INTO '
                  '    journeys(user, station_from, station_to, '
                  '    time_in, time_out, cost) '
                  'VALUES(?, ?, ?, ?, ?, ?)')

def find_user(cursor, username):
    q = 'SELECT id FROM users WHERE username=?'
    cursor.execute(q, (username, ))
    try:
        return int(cursor.fetchone()[0])
    except:
        raise Exception('No user with that username')

def add_to_db(cursor, uid, csv_path):
    for record in parse_csv(csv_path):
        if record['type'] == 'payment':
            cursor.execute(
                INSERT_PAYMENT,
                (uid, record['amount'], record['time'])
            )
        elif record['type'] == 'journey':
            cursor.execute(
                INSERT_JOURNEY,
                (uid, record['station_from'], record['station_to'],
                 record['time_in'], record['time_out'], record['cost'])
            )

class TellfLException(Exception):
    pass

class DB:

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

    def __enter__(self):
        self.conn = sqlite3.connect(os.path.join(self.location, DB_NAME))
        return self.conn.cursor()

    def __exit__(self, type, value, traceback):
        if traceback is None:
            self.conn.commit()
        self.conn.close()

def report(username, location, as_html):
    with DB(location) as cursor:
        uid = find_user(cursor, username)
        r = reports.create(cursor, uid)
        if as_html:
            print reports.asHtml(r)
        else:
            print reports.asText(r)

def add(username, csv_path, location):
    with DB(location) as cursor:
        uid = find_user(cursor, username)
        add_to_db(cursor, uid, csv_path)

def register(username, location):
    with DB(location) as cursor:
        try:
            q = 'INSERT INTO users(username) VALUES(?)'
            cursor.execute(q, (username, ))
        except:
            raise TellfLException('A user with that username already exists')

def install(location):
    os.makedirs(location)
    with DB(location) as cursor:
        sql = resource_string(__name__, 'assets/schema.sql')
        cursor.executescript(sql)
<|endoftext|>


<commit_before>"""Utilities and initialization code related to the Jinja2 templates
used by the benchmarking script.
"""

from jinja2 import Environment, FileSystemLoader

# Create the Template environment
env = Environment(
    loader=FileSystemLoader(['templates', 'templates/sites'])
)

#------ Below are Custom Filters used for Formatting --------

def is_blank(val):
    return (val is None) or (val == "nan")

def filter_blank(val):
    if is_blank(val):
        return ''
    else:
        return val

def filter_money(val):
    try:
        return "${:,}".format(val)
    except TypeError:
        return ''
    except ValueError:
        return val

def filter_money_accurate(val):
    try:
        return "${:,.3f}".format(val)
    except TypeError:
        return ''
    except ValueError:
        return val

def filter_number(val):
    try:
        return '{:,}'.format(val)
    except TypeError:
        return ''
    except ValueError:
        return val

def filter_percent(val):
    try:
        return "{:.0%}".format(val)
    except TypeError:
        return ''
    except ValueError:
        return val

# Assign the custom filters to the Environment
env.filters['blank'] = filter_blank
env.filters['money'] = filter_money
env.filters['money_accurate'] = filter_money_accurate
env.filters['number'] = filter_number
env.filters['percent'] = filter_percent

def get_template(template_filename):
    """Returns a template having the name 'template_file_name'
    """
    return env.get_template(template_filename)
<commit_msg>Updated template_util with newest custom filters.<commit_after>"""Utilities and initialization code related to the Jinja2 templates
used by the benchmarking script.
"""

from jinja2 import Environment, FileSystemLoader
import numpy as np

# Create the Template environment
env = Environment(
    loader=FileSystemLoader(['templates', 'templates/sites'])
)

#------ Below are Custom Filters used for Formatting --------

def is_blank(val):
    try:
        return not np.isfinite(val)
    except:
        return True

def filter_blank(val):
    if is_blank(val): return ''
    return val

def filter_money(val, precision = 0):
    if is_blank(val): return ''
    format_string = "${:,.%sf}" % precision
    return format_string.format(val)

def filter_number(val, precision = 0):
    if is_blank(val): return ''
    format_string = "{:,.%sf}" % precision
    return format_string.format(val)

def filter_percent(val, precision = 1):
    if is_blank(val): return ''
    format_string = "{:.%s%%}" % precision
    return format_string.format(val)

# -------------- Assign the custom filters to the Environment

env.filters['blank'] = filter_blank
env.filters['money'] = filter_money
env.filters['number'] = filter_number
env.filters['percent'] = filter_percent

# ------------- Utility Functions -----------------

def get_template(template_filename):
    """Returns a template having the name 'template_file_name'
    """
    return env.get_template(template_filename)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
@author: Federico Cerchiari <federicocerchiari@gmail.com>
Widgets
"""
from copy import copy
from itertools import zip_longest

import tempy.tags as tags
from .tools import AdjustableList
from .exceptions import WidgetDataError, WidgetError
from .tags import Html, Table, Tbody, Thead, Caption
from .tags import Tfoot, Td, Tr, Th, Li, Ul, Dl, Dd, Dt

class TempyPage(Html):
    """HTML Page widget.
    Builds an empty html page with some named common tags:
    - Doctype
    - Head - Title - description meta - keywords meta
    - Body
    Those elements are accessible as page attribute:
    >>> TempyPage.head.charset
    >>> TempyPage.body
    >>> TempyPage.head.title
    Provides an API to manage common meta tags directly.
    """
    __tag = Html._Html__tag

    def __init__(self, title=None, content=None, charset='UTF-8',
                 keywords=None, doctype=None, **kwargs):
        super().__init__(**kwargs)
        self.set_charset(charset)
        self.set_description(content)
        self.set_keywords(keywords or [])

    def init(self):
        self(
            head=tags.Head()(
                charset=tags.Meta(),
                description=tags.Meta(name='description'),
                keywords=tags.Meta(name='keywords'),
                title=tags.Title(),
            ),
            body=tags.Body()
        )

    def set_doctype(self, doctype):
        """Changes the <meta> charset tag (default charset in init is UTF-8)."""
        self.doctype.type_code = doctype
        return self

    def set_charset(self, charset):
        """Changes the <meta> charset tag (default charset in init is UTF-8)."""
        self.head.charset.attr(charset=charset)
        return self

    def set_description(self, description):
        """Changes the <meta> description tag."""
        self.head.description.attr(content=description)
        return self

    def set_keywords(self, keywords):
        """Changes the <meta> keywords tag."""
        self.head.keywords.attr(content=', '.join(keywords))
        return self

class TempyTable(Table):
    """Table widget.
    Creates a simple table structure using the give data, ora an empty table of the given size.self
    params:
    data: an iterable of iterables in this form [[col1, col2, col3], [col1, col2, col3]]
    rows, columns: size of the table if no data is given
    head: if True adds the table header using the first data row
    foot: if True add a table footer using the last data row
    caption: adds the caption to the table
    """
    __tag = Table._Table__tag

    def __init__(self, rows=0, cols=0, data=None, caption=None,
                 head=False, foot=False, **kwargs):
        super().__init__(**kwargs)
        self(body=Tbody())
        # Initialize empty datastructure if data is not given
        if not data:
            data = [[None for _ in range(cols)]
                    for _ in range(rows + sum((head, foot)))]
        if caption:
            self.make_caption(caption)
        if head:
            self.make_header(data.pop(0))
        if foot:
            self.make_footer(data.pop())
        if data:
            self.populate(data, resize_x=True)

    def _check_row_size(self, row):
        try:
            row_lenght = len(row)
        except TypeError:
            row_lenght = row
        if self.body.childs and max(map(len, self.body)) < row_lenght:
            raise WidgetDataError(self, 'The given data have more columns than the table.')

    def populate(self, data, resize_x=True, normalize=True):
        """Adds/Replace data in the table.
        data: an iterable of iterables in the form [[col1, col2, col3], [col1, col2, col3]]
        resize_x: if True, changes the x-size of the table according to the given data.
            If False and data have dimensions different from the existing table structure a WidgetDataError is raised.
        normalize: if True all the rows will have the same number of columns, if False, data structure is followed.
        """
        if data is None:
            raise WidgetDataError(self, 'Parameter data should be non-None, to empty the table use TempyTable.clear() or pass an empty list.')
        data = copy(data)

        if not self.body:
            # Table is empty
            self(body=Tbody())
        self.clear()

        max_data_x = max(map(len, data))
        if not resize_x:
            self._check_row_size(max_data_x)

        for t_row, d_row in zip_longest(self.body, data):
            if not d_row:
                t_row.remove()
            else:
                if not t_row:
                    t_row = Tr().append_to(self.body)
                if normalize:
                    d_row = AdjustableList(d_row).ljust(max_data_x, None)
                for t_cell, d_cell in zip_longest(t_row, d_row):
                    if not t_cell and resize_x:
                        t_cell = Td().append_to(t_row)
                    t_cell.empty()
                    if d_cell is not None:
                        t_cell(d_cell)
        return self

    def clear(self):
        return self.body.empty()

    def add_row(self, row_data, resize_x=True):
        """Adds a row at the end of the table"""
        if not resize_x:
            self._check_row_size(row_data)
        self.body(Tr()(Td()(cell) for cell in row_data))
        return self

    def pop_row(self, idr=None, tags=False):
        """Pops a row, default the last"""
        idr = idr if idr is not None else len(self.body) - 1
        row = self.body.pop(idr)
        return row if tags else [cell.childs[0] for cell in row]

    def pop_cell(self, idy=None, idx=None, tags=False):
        """Pops a cell, default the last of the last row"""
        idy = idy if idy is not None else len(self.body) - 1
        idx = idx if idx is not None else len(self.body[idy]) - 1
        cell = self.body[idy].pop(idx)
        return cell if tags else cell.childs[0]

    def _make_table_part(self, part, data):
        part_tag, inner_tag = {
            'header': (Thead, Th),
            'footer': (Tfoot, Td)
            }.get(part)
        part_instance = part_tag().append_to(self)
        if not hasattr(self, part):
            setattr(self, part, part_instance)
        return part_instance(Tr()(inner_tag()(col) for col in data))

    def make_header(self, head):
        """Makes the header row from the given data."""
        self._make_table_part('header', head)

    def make_footer(self, footer):
        """Makes the footer row from the given data."""
        self._make_table_part('footer', footer)

    def make_caption(self, caption):
        """Adds/Sobstitute the table's caption."""
        if not hasattr(self, 'caption'):
            self(caption=Caption())
        return self.caption.empty()(caption)

class TempyListMeta:
    """Widget for lists, manages the automatic generation starting from iterables and dicts.
    Plain iterables will produce plain lists, nested dicts will produce nested lists.
    Default list type is unordered. List type (ol, ul or dl) can be passed as
    string argument ("Ul", "Ol", "Dl"),
    list tag classes (tempy.tags.Ul, tempy.tags.Ol) or "_typ" special key in the given structure:
    >>> ul = TempyList()
    >>> ol = TempyList(typ='Ol')
    >>> ol = TempyList(typ=tempy.tags.Ol)
    >>> dl = TempyList(typ=tempy.tags.Dl)
    >>> ol = TempyList(struct={'_typ': tempy.tags.Ol})
    >>> ol = TempyList(struct={'_typ': 'Ol'})
    List building is made through the TempyList.populate method; this will trasform a python datastructure
    (list/tuple/dict/etc..) in a Tempy tree.
    """
    def __init__(self, struct=None):
        self._typ = self.__class__.__bases__[1]
        self._TempyList__tag = getattr(self._typ, '_%s__tag' % self._typ.__name__)
        super().__init__()
        self.populate(struct=struct)

    def populate(self, struct):
        """Generates the list tree.
        struct: if a list/set/tuple is given, a flat list is generated
        <*l>
v1

v2
...</*l>
        If the list type is 'Dl' a flat list without definitions is generated
        <*l>
v1

v2
...</*l>

        If the given struct is a dict, key contaninct lists/tuples/sets/dicts
        will be transformed in nested lists, and so on recursively, using dict
        keys as list items, and dict values as sublists. If type is 'Dl' each
        value will be transformed in definition (or list of definitions)
        except others dict. In that case, it will be transformed in <dfn> tags.

        >>> struct = {'ele1': None, 'ele2': ['sub21', 'sub22'], 'ele3': {'sub31': None, 'sub32': None, '_typ': 'Ol'}}
        >>> TempyList(struct=struct)

ele1

ele2

sub21

sub22

ele3

sub31

sub32

        """

        if struct is None:
            # Maybe raise? Empty the list?
            return self

        if isinstance(struct, (list, set, tuple)):
            struct = dict(zip_longest(struct, [None]))
        if not isinstance(struct, dict):
            raise WidgetDataError(self, 'List Input not managed, expected (dict, list), got %s' % type(struct))
        else:
           self.__process_struct(struct)
        return self

    def __process_struct(self, struct):
        for k, submenu in struct.items():
            item = Dt()(k) if self._typ == Dl else Li()(k)
            self(item)
            if submenu:
                if self._typ == Dl:
                    if isinstance(submenu, (list, set, tuple, dict)):
                        for e in submenu:
                            self(Dd()(e))
                    else:
                        self(Dd()(submenu))
                else:
                    item(TempyList(typ=self._typ, struct=submenu))

class TempyList:
    """TempyList is a class factory, it works for both ul and ol lists (TODO: dl).
    See TempyListMeta for TempyList methods and docstings."""
    def __new__(cls, typ=None, struct=None):
        try:
            typ = struct.pop('_typ')
        except (TypeError, KeyError, AttributeError):
            pass
        if isinstance(typ, str):
            try:
                typ = getattr(tags, typ)
            except AttributeError:
                raise WidgetError(cls, 'TempyList type not expected.')
        typ = typ or Ul
        cls_typ = type("TempyList%s" % typ.__name__, (TempyListMeta, typ), {})
        return cls_typ(struct=struct)
<commit_msg>Fixed code complexity<commit_after># -*- coding: utf-8 -*-
"""
@author: Federico Cerchiari <federicocerchiari@gmail.com>
Widgets
"""
from copy import copy
from itertools import zip_longest

import tempy.tags as tags
from .tools import AdjustableList
from .exceptions import WidgetDataError, WidgetError
from .tags import Html, Table, Tbody, Thead, Caption
from .tags import Tfoot, Td, Tr, Th, Li, Ul, Dl, Dd, Dt

class TempyPage(Html):
    """HTML Page widget.
    Builds an empty html page with some named common tags:
    - Doctype
    - Head - Title - description meta - keywords meta
    - Body
    Those elements are accessible as page attribute:
    >>> TempyPage.head.charset
    >>> TempyPage.body
    >>> TempyPage.head.title
    Provides an API to manage common meta tags directly.
    """
    __tag = Html._Html__tag

    def __init__(self, title=None, content=None, charset='UTF-8',
                 keywords=None, doctype=None, **kwargs):
        super().__init__(**kwargs)
        self.set_charset(charset)
        self.set_description(content)
        self.set_keywords(keywords or [])

    def init(self):
        self(
            head=tags.Head()(
                charset=tags.Meta(),
                description=tags.Meta(name='description'),
                keywords=tags.Meta(name='keywords'),
                title=tags.Title(),
            ),
            body=tags.Body()
        )

    def set_doctype(self, doctype):
        """Changes the <meta> charset tag (default charset in init is UTF-8)."""
        self.doctype.type_code = doctype
        return self

    def set_charset(self, charset):
        """Changes the <meta> charset tag (default charset in init is UTF-8)."""
        self.head.charset.attr(charset=charset)
        return self

    def set_description(self, description):
        """Changes the <meta> description tag."""
        self.head.description.attr(content=description)
        return self

    def set_keywords(self, keywords):
        """Changes the <meta> keywords tag."""
        self.head.keywords.attr(content=', '.join(keywords))
        return self

class TempyTable(Table):
    """Table widget.
    Creates a simple table structure using the give data, ora an empty table of the given size.self
    params:
    data: an iterable of iterables in this form [[col1, col2, col3], [col1, col2, col3]]
    rows, columns: size of the table if no data is given
    head: if True adds the table header using the first data row
    foot: if True add a table footer using the last data row
    caption: adds the caption to the table
    """
    __tag = Table._Table__tag

    def __init__(self, rows=0, cols=0, data=None, caption=None,
                 head=False, foot=False, **kwargs):
        super().__init__(**kwargs)
        self(body=Tbody())
        # Initialize empty datastructure if data is not given
        if not data:
            data = [[None for _ in range(cols)]
                    for _ in range(rows + sum((head, foot)))]
        if caption:
            self.make_caption(caption)
        if head:
            self.make_header(data.pop(0))
        if foot:
            self.make_footer(data.pop())
        if data:
            self.populate(data, resize_x=True)

    def _check_row_size(self, row):
        try:
            row_lenght = len(row)
        except TypeError:
            row_lenght = row
        if self.body.childs and max(map(len, self.body)) < row_lenght:
            raise WidgetDataError(self, 'The given data have more columns than the table.')

    def populate(self, data, resize_x=True, normalize=True):
        """Adds/Replace data in the table.
        data: an iterable of iterables in the form [[col1, col2, col3], [col1, col2, col3]]
        resize_x: if True, changes the x-size of the table according to the given data.
            If False and data have dimensions different from the existing table structure a WidgetDataError is raised.
        normalize: if True all the rows will have the same number of columns, if False, data structure is followed.
        """
        if data is None:
            raise WidgetDataError(self, 'Parameter data should be non-None, to empty the table use TempyTable.clear() or pass an empty list.')
        data = copy(data)

        if not self.body:
            # Table is empty
            self(body=Tbody())
        self.clear()

        max_data_x = max(map(len, data))
        if not resize_x:
            self._check_row_size(max_data_x)

        for t_row, d_row in zip_longest(self.body, data):
            if not d_row:
                t_row.remove()
            else:
                if not t_row:
                    t_row = Tr().append_to(self.body)
                if normalize:
                    d_row = AdjustableList(d_row).ljust(max_data_x, None)
                for t_cell, d_cell in zip_longest(t_row, d_row):
                    if not t_cell and resize_x:
                        t_cell = Td().append_to(t_row)
                    t_cell.empty()
                    if d_cell is not None:
                        t_cell(d_cell)
        return self

    def clear(self):
        return self.body.empty()

    def add_row(self, row_data, resize_x=True):
        """Adds a row at the end of the table"""
        if not resize_x:
            self._check_row_size(row_data)
        self.body(Tr()(Td()(cell) for cell in row_data))
        return self

    def pop_row(self, idr=None, tags=False):
        """Pops a row, default the last"""
        idr = idr if idr is not None else len(self.body) - 1
        row = self.body.pop(idr)
        return row if tags else [cell.childs[0] for cell in row]

    def pop_cell(self, idy=None, idx=None, tags=False):
        """Pops a cell, default the last of the last row"""
        idy = idy if idy is not None else len(self.body) - 1
        idx = idx if idx is not None else len(self.body[idy]) - 1
        cell = self.body[idy].pop(idx)
        return cell if tags else cell.childs[0]

    def _make_table_part(self, part, data):
        part_tag, inner_tag = {
            'header': (Thead, Th),
            'footer': (Tfoot, Td)
            }.get(part)
        part_instance = part_tag().append_to(self)
        if not hasattr(self, part):
            setattr(self, part, part_instance)
        return part_instance(Tr()(inner_tag()(col) for col in data))

    def make_header(self, head):
        """Makes the header row from the given data."""
        self._make_table_part('header', head)

    def make_footer(self, footer):
        """Makes the footer row from the given data."""
        self._make_table_part('footer', footer)

    def make_caption(self, caption):
        """Adds/Sobstitute the table's caption."""
        if not hasattr(self, 'caption'):
            self(caption=Caption())
        return self.caption.empty()(caption)

class TempyListMeta:
    """Widget for lists, manages the automatic generation starting from iterables and dicts.
    Plain iterables will produce plain lists, nested dicts will produce nested lists.
    Default list type is unordered. List type (ol, ul or dl) can be passed as
    string argument ("Ul", "Ol", "Dl"),
    list tag classes (tempy.tags.Ul, tempy.tags.Ol) or "_typ" special key in the given structure:
    >>> ul = TempyList()
    >>> ol = TempyList(typ='Ol')
    >>> ol = TempyList(typ=tempy.tags.Ol)
    >>> dl = TempyList(typ=tempy.tags.Dl)
    >>> ol = TempyList(struct={'_typ': tempy.tags.Ol})
    >>> ol = TempyList(struct={'_typ': 'Ol'})
    List building is made through the TempyList.populate method; this will trasform a python datastructure
    (list/tuple/dict/etc..) in a Tempy tree.
    """
    def __init__(self, struct=None):
        self._typ = self.__class__.__bases__[1]
        self._TempyList__tag = getattr(self._typ, '_%s__tag' % self._typ.__name__)
        super().__init__()
        self.populate(struct=struct)

    def populate(self, struct):
        """Generates the list tree.
        struct: if a list/set/tuple is given, a flat list is generated
        <*l>
v1

v2
...</*l>
        If the list type is 'Dl' a flat list without definitions is generated
        <*l>
v1

v2
...</*l>

        If the given struct is a dict, key contaninct lists/tuples/sets/dicts
        will be transformed in nested lists, and so on recursively, using dict
        keys as list items, and dict values as sublists. If type is 'Dl' each
        value will be transformed in definition (or list of definitions)
        except others dict. In that case, it will be transformed in <dfn> tags.

        >>> struct = {'ele1': None, 'ele2': ['sub21', 'sub22'], 'ele3': {'sub31': None, 'sub32': None, '_typ': 'Ol'}}
        >>> TempyList(struct=struct)

ele1

ele2

sub21

sub22

ele3

sub31

sub32

        """

        if struct is None:
            # Maybe raise? Empty the list?
            return self

        if isinstance(struct, (list, set, tuple)):
            struct = dict(zip_longest(struct, [None]))
        if not isinstance(struct, dict):
            raise WidgetDataError(self, 'List Input not managed, expected (dict, list), got %s' % type(struct))
        else:
            if self._typ == Dl:
                self.__process_dl_struct(struct)
            else:
                self.__process_li_struct(struct)

        return self

    def __process_li_struct(self, struct):
        for k, submenu in struct.items():
            item = Li()(k)
            self(item)
            if submenu:
                item(TempyList(typ=self._typ, struct=submenu))

    def __process_dl_struct(self, struct):
        for k, submenu in struct.items():
            item = Dt()(k)
            self(item)
            if submenu:
                if isinstance(submenu, (list, set, tuple, dict)):
                    for elem in submenu:
                        self(Dd()(elem))
                else:
                    self(Dd()(submenu))

class TempyList:
    """TempyList is a class factory, it works for both ul and ol lists (TODO: dl).
    See TempyListMeta for TempyList methods and docstings."""
    def __new__(cls, typ=None, struct=None):
        try:
            typ = struct.pop('_typ')
        except (TypeError, KeyError, AttributeError):
            pass
        if isinstance(typ, str):
            try:
                typ = getattr(tags, typ)
            except AttributeError:
                raise WidgetError(cls, 'TempyList type not expected.')
        typ = typ or Ul
        cls_typ = type("TempyList%s" % typ.__name__, (TempyListMeta, typ), {})
        return cls_typ(struct=struct)
<|endoftext|>


<commit_before>
<commit_msg>Broken testing<commit_after>from test_ghx_aggregated_load import *
from test_ghx_base import *
from test_ghx_constants import *
from test_ghx_fluids import *
from test_ghxArray_Euler import *
from test_ghxArray_Lagrange import *
<|endoftext|>


<commit_before>import datetime

class NewDate(datetime.date):
    @classmethod
    def today(cls):
        return cls(2012, 12, 10)

class DeferredHelper(object):
    def __init__(self, fire_callback=True, fire_errback=False, data=None):
        self.callbacks = []
        self.errbacks = []
        self.fire_callback = fire_callback
        self.fire_errback = fire_errback
        self.data = data
        self.args = None

    def _fire(self):
        def _walk_func_list(func_list):
            res = None
            for callback in func_list:
                res = callback(self.data)
                if res is None:
                    break

        if self.fire_callback and len(self.callbacks) > 0:
            _walk_func_list(self.callbacks)

        if self.fire_errback and len(self.errbacks) > 0:
            _walk_func_list(self.errbacks)

    def addCallback(self, callback):
        self.callbacks.append(callback)

    def addErrback(self, errback):
        self.errbacks.append(errback)

    def addCallbacks(self, callback, errback):
        self.callbacks.append(callback)
        self.errbacks.append(errback)

    def __call__(self, *args):
        self.args = args
        return self
<commit_msg>test._helpers: add mock_is_admin<commit_after>import datetime
import mock

from functools import wraps

class NewDate(datetime.date):
    @classmethod
    def today(cls):
        return cls(2012, 12, 10)

class DeferredHelper(object):
    def __init__(self, fire_callback=True, fire_errback=False, data=None):
        self.callbacks = []
        self.errbacks = []
        self.fire_callback = fire_callback
        self.fire_errback = fire_errback
        self.data = data
        self.args = None

    def _fire(self):
        def _walk_func_list(func_list):
            res = None
            for callback in func_list:
                res = callback(self.data)
                if res is None:
                    break

        if self.fire_callback and len(self.callbacks) > 0:
            _walk_func_list(self.callbacks)

        if self.fire_errback and len(self.errbacks) > 0:
            _walk_func_list(self.errbacks)

    def addCallback(self, callback):
        self.callbacks.append(callback)

    def addErrback(self, errback):
        self.errbacks.append(errback)

    def addCallbacks(self, callback, errback):
        self.callbacks.append(callback)
        self.errbacks.append(errback)

    def __call__(self, *args):
        self.args = args
        return self

def mock_is_admin(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        with mock.patch("lala.pluginmanager.is_admin") as mocked:
            mocked.return_value = True
            return f(*args, **kwargs)
    return wrapper
<|endoftext|>