<commit_before># -*- coding: utf-8 -*-
###############################################################################
#
#    Copyright (C) 2001-2014 Micronaet SRL (<http://www.micronaet.it>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################
import sys
import os
import shutil
import ConfigParser
from datetime import datetime

# -----------------------------------------------------------------------------
# Utility:
# -----------------------------------------------------------------------------
def clean_text(text, length, uppercase=False, error=None, truncate=False):
    ''' Return clean text with limit cut
        Log in error if over length
    '''
    if error is None:
        error = []

    text = text.strip()
    if len(text) > length:
        if truncate:
            return text = text[:length]
        else:
            error.append('Text: %s > %s' % (text, length))
    if uppercase:
        return text.upper()
    return text

def clean_date(italian_date, separator='', out_format='iso', error=None):
    ''' Return clean text with limit cut
        Log in error if over length
    '''
    if error is None:
        error = []
    italian_date = italian_date.split(' ')[0] # remove hour block
    if len(italian_date) != 10:
        error.append('Error not italian date: %s' % italian_date)
        # not stopped
    if out_format == 'iso':
        return '%s%s%s%s%s' % (
            italian_date[-4:],
            separator,
            italian_date[3:5],
            separator,
            italian_date[:2],
            )
    elif out_format == 'italian':
        return '%s%s%s%s%s' % (
            italian_date[:2],
            separator,
            italian_date[3:5],
            separator,
            italian_date[-4:],
            )
    elif out_format == 'english':
        return '%s%s%s%s%s' % (
            italian_date[3:5],
            separator,
            italian_date[:2],
            separator,
            italian_date[-4:],
            )
    else: # incorrect format:
        return italian_date # nothing todo

def clean_float(value, length, decimal=3, separator='.', error=None):
    ''' Clean float and return float format
    '''
    if error is None:
        error = []
    try:
        value = value.replace(',', '.')
        float_value = float(value.strip())
    except:
        error.append('Not a float: %s' % value)
        float_value = 0.0

    mask = '%%%s.%sf' % (length, decimal)
    res = mask % float_value
    res = res.replace('.', separator)
    return res

def log_on_file(message, mode='INFO', file_list=None, verbose=True):
    ''' Write message in file list passed
        mode: INFO WARNING ERROR
    '''
    if file_list is None:
        return False

    message_log = '%s. [%s] %s\n' % (datetime.now(), mode, message)
    if verbose:
        print message_log.strip()

    for f_log in file_list:
        if not f_log:
            continue
        f_log.write(message_log)
    return True

# -----------------------------------------------------------------------------
# Read configuration parameter:
# -----------------------------------------------------------------------------
try:
    company = sys.argv[1]
    if len(company) > 3:
        print 'Run: python ./pre_process.py CMP (must be max 3 char length)'
        sys.exit()
except:
    print 'Launch program as:\n\npython ./pre_process.py CMP'

try:
    cfg_file = os.path.expanduser('./%s.cfg' % company)
except:
    print 'Config file not found: %s.cfg' % company

config = ConfigParser.ConfigParser()
config.read([cfg_file])

# In parameters:
in_check = os.path.expanduser(config.get('in', 'check'))
in_path = os.path.expanduser(config.get('in', 'path'))
in_history = os.path.expanduser(config.get('in', 'history'))
in_log = os.path.expanduser(config.get('in', 'log'))
in_schedule = os.path.expanduser(config.get('in', 'schedule'))

# Outp arameters:
out_check = os.path.expanduser(config.get('out', 'check'))
out_path = os.path.expanduser(config.get('out', 'path'))
out_log = os.path.expanduser(config.get('out', 'log'))
out_schedule = os.path.expanduser(config.get('out', 'schedule'))
out_original = os.path.expanduser(config.get('out', 'original'))

# Mail parameters:
mail_error = config.get('mail', 'error')
mail_info = config.get('mail', 'info')

# TODO put in config file:
separator = config.get('file', 'separator')
tot_col = eval(config.get('file', 'tot_col'))

# Calculated parameters:
f_in_schedule = False # open(in_schedule, 'a')
f_in_log = False #open(in_log, 'a')

f_out_schedule = open(out_schedule, 'a')
f_out_log = open(out_log, 'a')

# -----------------------------------------------------------------------------
# Check mount folder:
# -----------------------------------------------------------------------------
if in_check:
    if not os.path.exists(in_check):
        log_on_file(
            'Cannot mount IN folder: %s' % in_path, mode='ERROR', file_list=[
                f_in_schedule, f_out_schedule])
        sys.exit()

if out_check:
    if not os.path.exists(out_check):
        log_on_file(
            'Cannot mount OUT folder: %s' % out_path, mode='ERROR', file_list=[
                f_in_schedule, f_out_schedule])
        sys.exit()

# -----------------------------------------------------------------------------
# Read IN folder:
# -----------------------------------------------------------------------------
log_on_file(
    'Start import order mode: %s' % company, mode='INFO', file_list=[
        f_in_schedule, f_out_schedule])

for root, dirs, files in os.walk(in_path):
    log_on_file(
        'Read root folder: %s [%s]' % (root, company),
        mode='INFO',
        file_list=[f_in_log, f_out_log])

    for f in files:
        # Fullname needed:
        file_in = os.path.join(root, f)
        file_history = os.path.join(in_history, f)
        file_original = os.path.join(out_original, f)

        create_date = datetime.fromtimestamp(
            os.path.getctime(file_in)
            ).strftime('%Y%m%d_%H%M%S')

        file_out = os.path.join(
            out_path, '%s' % (
                '%s_%s' % (create_date, f),
                )) # TODO change name

        # ---------------------------------------------------------------------
        # Read input file:
        # ---------------------------------------------------------------------
        error = []
        f_in = open(file_in, 'r')
        row_out = []
        log_on_file(
            'Parsing file: %s [%s]' % (file_in, company),
            mode='INFO',
            file_list=[f_in_log, f_out_log])
        for line in f_in:
            row = line.split(separator)
            if len(row) != tot_col:
                log_on_file(
                    'Different number columns: %s [%s]' % (file_in, company),
                    mode='INFO', file_list=[f_in_log, f_out_log])
                break

            # -----------------------------------------------------------------
            # Read fields:
            # -----------------------------------------------------------------
            date = create_date[:8] # From create date

            deadline = clean_date(row[0], separator='', out_format='iso',
                error=error)

            # Char:
            code = clean_text(row[1], 9, error=error)
            cook_center = clean_text(row[2], 110, error=error)[:60]# XXX Trunc
            address = clean_text(row[3], 60, error=error)
            cap = clean_text(row[4], 5, error=error) # XXX Check
            city = clean_text(row[5], 60, error=error) # XXX Check
            province = clean_text(row[6], 4, uppercase=True, error=error)
            order = clean_text(row[7], 10, error=error)
            default_code = clean_text(row[8], 16, uppercase=True, error=error)
            default_code_supplier = clean_text(
                row[9], 16, uppercase=True, error=error)
            name = clean_text(row[10], 60, error=error, truncate=True)
            um = clean_text(row[11], 2, uppercase=True, error=error)

            # Float:
            quantity = clean_float(row[12], 15, 2, error=error)
            price = clean_float(row[13], 15, 3, error=error)
            cost = clean_float(row[14], 15, 3, error=error)

            # -----------------------------------------------------------------
            # Convert row input file:
            # -----------------------------------------------------------------
            row_out.append(
                '%3s|%-8s|%-9s|%-60s|%-60s|%-5s|%-60s|%-4s|%-10s|%-16s|%-16s|'
                '%-60s|%-2s|%-15s|%-15s|%-15s|%-8s\r\n' % (
                    company, # depend on parameter
                    deadline,
                    code,
                    cook_center,
                    address,
                    cap,
                    city,
                    province,
                    order,
                    default_code,
                    default_code_supplier,
                    name,
                    um,
                    quantity,
                    price,
                    cost,
                    date,
                    ))
        f_in.close()

        if error:
            log_on_file(
                'Error reading file: %s [%s]' % (error, company),
                mode='ERROR', file_list=[f_in_log, f_out_log])
            continue

        # ---------------------------------------------------------------------
        # Save file out:
        # ---------------------------------------------------------------------
        error = []
        # Convert file:
        f_out = open(file_out, 'w')

        for line in row_out:
            f_out.write(line)
        f_out.close()

        if error:
            log_on_file(
                'Error generating out file: %s [%s]' % (error, company),
                mode='ERROR', file_list=[f_in_log, f_out_log])
            continue
        else:
            log_on_file(
                'Created file: %s [%s]' % (file_out, company),
                mode='INFO', file_list=[f_in_log, f_out_log])

        # ---------------------------------------------------------------------
        # History the in file:
        # ---------------------------------------------------------------------
        error = ''
        try:
            os.remove(file_original)
        except:
            pass
        try:
            shutil.copy(file_in, file_original)
            log_on_file(
                'History company file: %s [%s]' % (file_original, company),
                mode='INFO', file_list=[f_in_log, f_out_log])
        except:
            error += 'Error copy orginal folder %s >> %s' % (
                file_in, file_original)
        try:
            os.remove(file_history)
        except:
            pass
        try:
            shutil.move(file_in, file_history)
            log_on_file(
                'History supplier file: %s [%s]' % (file_history, company),
                mode='INFO', file_list=[f_in_log, f_out_log])
        except:
            error += 'Error moving file %s >> %s' % (file_in, file_history)

        if error:
            log_on_file(
                error, mode='ERROR', file_list=[f_in_log, f_out_log])
            # XXX TODO delete OUT file
            continue
    break # only root folder is read

log_on_file(
    'End import order mode: %s' % company, mode='INFO', file_list=[
        f_in_schedule, f_out_schedule])

# -----------------------------------------------------------------------------
# History the in file:
# -----------------------------------------------------------------------------
try:
    if f_out_log:
        f_out_log.close()
    if f_in_schedule:
        f_in_schedule.close()
    if f_in_log:
        f_in_log.close()
    if f_out_schedule:
        f_out_schedule.close()
except:
    print 'Error closing log file'
<commit_msg>generic bugfix (fast message)<commit_after># -*- coding: utf-8 -*-
###############################################################################
#
#    Copyright (C) 2001-2014 Micronaet SRL (<http://www.micronaet.it>).
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU Affero General Public License as published
#    by the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU Affero General Public License for more details.
#
#    You should have received a copy of the GNU Affero General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################
import sys
import os
import shutil
import ConfigParser
from datetime import datetime

# -----------------------------------------------------------------------------
# Utility:
# -----------------------------------------------------------------------------
def clean_text(text, length, uppercase=False, error=None, truncate=False):
    ''' Return clean text with limit cut
        Log in error if over length
    '''
    if error is None:
        error = []

    text = text.strip()
    if len(text) > length:
        if truncate:
            text = text[:length]
        else:
            error.append('Text: %s > %s' % (text, length))
    if uppercase:
        return text.upper()
    return text

def clean_date(italian_date, separator='', out_format='iso', error=None):
    ''' Return clean text with limit cut
        Log in error if over length
    '''
    if error is None:
        error = []
    italian_date = italian_date.split(' ')[0] # remove hour block
    if len(italian_date) != 10:
        error.append('Error not italian date: %s' % italian_date)
        # not stopped
    if out_format == 'iso':
        return '%s%s%s%s%s' % (
            italian_date[-4:],
            separator,
            italian_date[3:5],
            separator,
            italian_date[:2],
            )
    elif out_format == 'italian':
        return '%s%s%s%s%s' % (
            italian_date[:2],
            separator,
            italian_date[3:5],
            separator,
            italian_date[-4:],
            )
    elif out_format == 'english':
        return '%s%s%s%s%s' % (
            italian_date[3:5],
            separator,
            italian_date[:2],
            separator,
            italian_date[-4:],
            )
    else: # incorrect format:
        return italian_date # nothing todo

def clean_float(value, length, decimal=3, separator='.', error=None):
    ''' Clean float and return float format
    '''
    if error is None:
        error = []
    try:
        value = value.replace(',', '.')
        float_value = float(value.strip())
    except:
        error.append('Not a float: %s' % value)
        float_value = 0.0

    mask = '%%%s.%sf' % (length, decimal)
    res = mask % float_value
    res = res.replace('.', separator)
    return res

def log_on_file(message, mode='INFO', file_list=None, verbose=True):
    ''' Write message in file list passed
        mode: INFO WARNING ERROR
    '''
    if file_list is None:
        return False

    message_log = '%s. [%s] %s\n' % (datetime.now(), mode, message)
    if verbose:
        print message_log.strip()

    for f_log in file_list:
        if not f_log:
            continue
        f_log.write(message_log)
    return True

# -----------------------------------------------------------------------------
# Read configuration parameter:
# -----------------------------------------------------------------------------
try:
    company = sys.argv[1]
    if len(company) > 3:
        print 'Run: python ./pre_process.py CMP (must be max 3 char length)'
        sys.exit()
except:
    print 'Launch program as:\n\npython ./pre_process.py CMP'

try:
    cfg_file = os.path.expanduser('./%s.cfg' % company)
except:
    print 'Config file not found: %s.cfg' % company

config = ConfigParser.ConfigParser()
config.read([cfg_file])

# In parameters:
in_check = os.path.expanduser(config.get('in', 'check'))
in_path = os.path.expanduser(config.get('in', 'path'))
in_history = os.path.expanduser(config.get('in', 'history'))
in_log = os.path.expanduser(config.get('in', 'log'))
in_schedule = os.path.expanduser(config.get('in', 'schedule'))

# Outp arameters:
out_check = os.path.expanduser(config.get('out', 'check'))
out_path = os.path.expanduser(config.get('out', 'path'))
out_log = os.path.expanduser(config.get('out', 'log'))
out_schedule = os.path.expanduser(config.get('out', 'schedule'))
out_original = os.path.expanduser(config.get('out', 'original'))

# Mail parameters:
mail_error = config.get('mail', 'error')
mail_info = config.get('mail', 'info')

# TODO put in config file:
separator = config.get('file', 'separator')
tot_col = eval(config.get('file', 'tot_col'))

# Calculated parameters:
f_in_schedule = False # open(in_schedule, 'a')
f_in_log = False #open(in_log, 'a')

f_out_schedule = open(out_schedule, 'a')
f_out_log = open(out_log, 'a')

# -----------------------------------------------------------------------------
# Check mount folder:
# -----------------------------------------------------------------------------
if in_check:
    if not os.path.exists(in_check):
        log_on_file(
            'Cannot mount IN folder: %s' % in_path, mode='ERROR', file_list=[
                f_in_schedule, f_out_schedule])
        sys.exit()

if out_check:
    if not os.path.exists(out_check):
        log_on_file(
            'Cannot mount OUT folder: %s' % out_path, mode='ERROR', file_list=[
                f_in_schedule, f_out_schedule])
        sys.exit()

# -----------------------------------------------------------------------------
# Read IN folder:
# -----------------------------------------------------------------------------
log_on_file(
    'Start import order mode: %s' % company, mode='INFO', file_list=[
        f_in_schedule, f_out_schedule])

for root, dirs, files in os.walk(in_path):
    log_on_file(
        'Read root folder: %s [%s]' % (root, company),
        mode='INFO',
        file_list=[f_in_log, f_out_log])

    for f in files:
        # Fullname needed:
        file_in = os.path.join(root, f)
        file_history = os.path.join(in_history, f)
        file_original = os.path.join(out_original, f)

        create_date = datetime.fromtimestamp(
            os.path.getctime(file_in)
            ).strftime('%Y%m%d_%H%M%S')

        file_out = os.path.join(
            out_path, '%s' % (
                '%s_%s' % (create_date, f),
                )) # TODO change name

        # ---------------------------------------------------------------------
        # Read input file:
        # ---------------------------------------------------------------------
        error = []
        f_in = open(file_in, 'r')
        row_out = []
        log_on_file(
            'Parsing file: %s [%s]' % (file_in, company),
            mode='INFO',
            file_list=[f_in_log, f_out_log])
        for line in f_in:
            row = line.split(separator)
            if len(row) != tot_col:
                log_on_file(
                    'Different number columns: %s [%s]' % (file_in, company),
                    mode='INFO', file_list=[f_in_log, f_out_log])
                break

            # -----------------------------------------------------------------
            # Read fields:
            # -----------------------------------------------------------------
            date = create_date[:8] # From create date

            deadline = clean_date(row[0], separator='', out_format='iso',
                error=error)

            # Char:
            code = clean_text(row[1], 9, error=error)
            cook_center = clean_text(row[2], 110, error=error)[:60]# XXX Trunc
            address = clean_text(row[3], 60, error=error)
            cap = clean_text(row[4], 5, error=error) # XXX Check
            city = clean_text(row[5], 60, error=error) # XXX Check
            province = clean_text(row[6], 4, uppercase=True, error=error)
            order = clean_text(row[7], 10, error=error)
            default_code = clean_text(row[8], 16, uppercase=True, error=error)
            default_code_supplier = clean_text(
                row[9], 16, uppercase=True, error=error)
            name = clean_text(row[10], 60, error=error, truncate=True)
            um = clean_text(row[11], 2, uppercase=True, error=error)

            # Float:
            quantity = clean_float(row[12], 15, 2, error=error)
            price = clean_float(row[13], 15, 3, error=error)
            cost = clean_float(row[14], 15, 3, error=error)

            # -----------------------------------------------------------------
            # Convert row input file:
            # -----------------------------------------------------------------
            row_out.append(
                '%3s|%-8s|%-9s|%-60s|%-60s|%-5s|%-60s|%-4s|%-10s|%-16s|%-16s|'
                '%-60s|%-2s|%-15s|%-15s|%-15s|%-8s\r\n' % (
                    company, # depend on parameter
                    deadline,
                    code,
                    cook_center,
                    address,
                    cap,
                    city,
                    province,
                    order,
                    default_code,
                    default_code_supplier,
                    name,
                    um,
                    quantity,
                    price,
                    cost,
                    date,
                    ))
        f_in.close()

        if error:
            log_on_file(
                'Error reading file: %s [%s]' % (error, company),
                mode='ERROR', file_list=[f_in_log, f_out_log])
            continue

        # ---------------------------------------------------------------------
        # Save file out:
        # ---------------------------------------------------------------------
        error = []
        # Convert file:
        f_out = open(file_out, 'w')

        for line in row_out:
            f_out.write(line)
        f_out.close()

        if error:
            log_on_file(
                'Error generating out file: %s [%s]' % (error, company),
                mode='ERROR', file_list=[f_in_log, f_out_log])
            continue
        else:
            log_on_file(
                'Created file: %s [%s]' % (file_out, company),
                mode='INFO', file_list=[f_in_log, f_out_log])

        # ---------------------------------------------------------------------
        # History the in file:
        # ---------------------------------------------------------------------
        error = ''
        try:
            os.remove(file_original)
        except:
            pass
        try:
            shutil.copy(file_in, file_original)
            log_on_file(
                'History company file: %s [%s]' % (file_original, company),
                mode='INFO', file_list=[f_in_log, f_out_log])
        except:
            error += 'Error copy orginal folder %s >> %s' % (
                file_in, file_original)
        try:
            os.remove(file_history)
        except:
            pass
        try:
            shutil.move(file_in, file_history)
            log_on_file(
                'History supplier file: %s [%s]' % (file_history, company),
                mode='INFO', file_list=[f_in_log, f_out_log])
        except:
            error += 'Error moving file %s >> %s' % (file_in, file_history)

        if error:
            log_on_file(
                error, mode='ERROR', file_list=[f_in_log, f_out_log])
            # XXX TODO delete OUT file
            continue
    break # only root folder is read

log_on_file(
    'End import order mode: %s' % company, mode='INFO', file_list=[
        f_in_schedule, f_out_schedule])

# -----------------------------------------------------------------------------
# History the in file:
# -----------------------------------------------------------------------------
try:
    if f_out_log:
        f_out_log.close()
    if f_in_schedule:
        f_in_schedule.close()
    if f_in_log:
        f_in_log.close()
    if f_out_schedule:
        f_out_schedule.close()
except:
    print 'Error closing log file'
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
Given a BAM file and a genomic reference, for each position covered in the
BAM file, show the reference base, and the number of A/T/C/G's and InDels.

You can also optionally filter out all bases whose quality score is below a
threshold, bases that aren't covered by enough reads, bases that have no
variation compared to reference, or bases whose variation is too low.

The output is a tab-delimited file that contains the following for each base:

chromosome
position (1-based)
reference base
# reads that contain this base
Entropy
# A calls
# C calls
# G calls
# T calls
# deletions
# gaps
# inserts

Entropy is sum(a..t) { p log2 p } where p = freq(+pseudocount) / genomic freq.
pseudo count = genomic freq * sqrt(N)

We use the following genomic frequencies: A 0.3, C 0.2, G 0.2, T 0.3
"""

import os,sys,math,collections
from support.eta import ETA
import pysam

def usage():
    base = os.path.basename(sys.argv[0])
    print __doc__
    print """
Usage: %s {opts} in.bam {chrom:start-end}

Options:
-ref   val    Include reference basecalls from this file
-qual  val    Minimum quality level to use in calculations
              (numeric, Sanger scale) (default 0)

-count val    Report only bases with this minimum number of reads covering it
              (default 0)

-mask  val    The bitmask to use for filtering reads (default 1540)

""" % (base)
    sys.exit(1)

__genomic_freq = {'A': 0.3, 'C': 0.2, 'G': 0.2,'T': 0.3}

def calc_entropy(a,c,t,g):
    counts = {'A':a,'C':c,'G':g,'T':t,}

    N = counts['A'] + counts['C'] + counts['G'] + counts['T']
    if N == 0:
        return 0

    N_sqrt = math.sqrt(N)

    count_pseudo = {}
    N_pseudo = 0

    for base in 'ATCG':
        count_pseudo[base] = counts[base] + (__genomic_freq[base] * N_sqrt)
        N_pseudo += count_pseudo[base]

    acc = 0
    for base in 'ATCG':
        p = float(count_pseudo[base]) / N_pseudo / __genomic_freq[base]
        acc += (p * math.log(p,2))

    return acc

MappingRecord = collections.namedtuple('MappingRecord','qpos cigar_op base qual read')
MappingPos = collections.namedtuple('MappingPos','tid pos records')
BasePosition = collections.namedtuple('BasePosition','tid pos total a c g t n deletions gaps insertions reads')

class BamBaseCaller(object):
    def __init__(self, bam_fname, min_qual=0, min_count=0, chrom=None, start=0, end=0, mask=1540,quiet=False):
        self.bam = pysam.Samfile(bam_fname,'rb')
        self.min_qual = min_qual
        self.min_count = 0

        self.chrom = chrom
        self.start = start
        self.end = end

        self.mask = mask
        self.quiet = quiet

        if chrom and start and end:
            def _gen():
                for p in self.bam.fetch(chrom,start,end):
                    yield p
        else:
            def _gen():
                for p in self.bam:
                    yield p
        self._gen = _gen

        self.buffer = None
        self.current_tid = None

    def close(self):
        self.bam.close()

    def _calc_pos(self,tid,pos,records):
        counts = {'A':0,'C':0,'G':0,'T':0,'N':0}
        insertions = {}
        deletions = 0
        gaps = 0
        total = 0
        reads = []

        for qpos,cigar_op,base,qual,read in records:
            if cigar_op == 0: # M
                if qual >= self.min_qual and (read.flag & self.mask) == 0 :
                    total += 1
                    reads.append(read)

                    counts[base] += 1
            elif cigar_op == 1: # I
                if qual >= self.min_qual and (read.flag & self.mask) == 0 :
                    reads.append(read)

                    if not base in insertions:
                        insertions[base] = 1
                    else:
                        insertions[base] += 1
            elif cigar_op == 2: # D
                deletions += 1
                #total += 1 # not sure these should be included,
                           # samtools mpileup includes them
                           # IGV doesn't

                reads.append(read)
            elif cigar_op == 3: # N
                gaps += 1
                reads.append(read)

        if total >= self.min_count:
            return BasePosition(tid,pos,total,counts['A'],counts['C'],counts['G'],counts['T'],counts['N'],deletions,gaps,insertions,reads)

    def fetch(self):
        self.current_tid = None
        self.buffer = collections.deque()
        if not self.quiet:
            eta = ETA(0,bamfile=self.bam)
        else:
            eta = None

        for read in self._gen():
            if eta:
                eta.print_status(extra='%s:%s (%s)' % (self.bam.references[read.tid],read.pos,len(self.buffer)),bam_pos=(read.tid,read.pos))

            if self.current_tid != read.tid: # new chromosome
                while self.buffer:
                    tid,pos,records = self.buffer.popleft()
                    yield self._calc_pos(tid,pos,records)

                self.current_tid = read.tid

            # handle all positions that are 5' of the current one
            while self.buffer and read.pos > self.buffer[0].pos:
                tid,pos,records = self.buffer.popleft()
                yield self._calc_pos(tid,pos,records)

            self._push_read(read)

        # flush buffer for the end
        while self.buffer:
            tid,pos,records = self.buffer.popleft()
            yield self._calc_pos(tid,pos,records)

        if eta:
            eta.done()

    def _push_read(self,read):

        if not self.buffer:
            self.buffer.append(MappingPos(read.tid, read.pos, []))

        while self.buffer[-1].pos < read.aend:
            self.buffer.append(MappingPos(read.tid, self.buffer[-1].pos+1, []))

        buf_idx = 0
        while self.buffer[0].pos < read.pos:
            buf_idx += 1

        read_idx = 0
        for op,length in read.cigar:
            if op == 0:
                for i in xrange(length):
                    self.buffer[buf_idx].records.append(MappingRecord(read_idx,op,read.seq[read_idx],read.qual[read_idx],read))
                    buf_idx += 1
                    read_idx += 1
            elif op == 1:
                inseq = ''
                inqual = 0
                for i in xrange(length):
                    inseq += read.seq[read_idx]
                    inqual += ord(read.qual[read_idx])-33
                    read_idx += 1

                inqual = inqual / len(inseq)

                self.buffer[buf_idx].records.append(MappingRecord(read_idx,op,inseq,inqual,read))

            elif op == 2:
                for i in xrange(length):
                    self.buffer[buf_idx].records.append(MappingRecord(read_idx,op,None,None,read))
                    buf_idx += 1

            elif op == 3:
                for i in xrange(length):
                    self.buffer[buf_idx].records.append(MappingRecord(read_idx,op,None,None,read))
                    buf_idx += 1

def bam_basecall(bam_fname,ref_fname,min_qual=0, min_count=0, chrom=None,start=None,end=None,mask=1540,quiet = False):
    if ref_fname:
        ref = pysam.Fastafile(ref_fname)
    else:
        ref = None

    sys.stdout.write('chrom\tpos\tref\tcount\tave mappings\tentropy\tA\tC\tG\tT\tN\tDeletions\tGaps\tInsertions\tInserts\n')
    bbc = BamBaseCaller(bam_fname,min_qual,min_count,chrom,start,end,mask,quiet)
    for basepos in bbc.fetch():
        if start and end:
            if basepos.pos < start or basepos.pos >= end:
                continue

        big_total = basepos.total + basepos.deletions + basepos.gaps + len(basepos.insertions)

        if big_total < min_count:
            continue

        if ref:
            refbase = ref.fetch(bbc.bam.references[basepos.tid],basepos.pos,basepos.pos+1).upper()
        else:
            refbase = 'N'

        entropy = calc_entropy(basepos.a,basepos.c,basepos.g,basepos.t)

        read_ih_acc = 0
        for read in basepos.reads:
            read_ih_acc += int(read.opt('IH'))

        inserts = []
        for insert in basepos.insertions:
            inserts.append((basepos.insertions[insert],insert))
        inserts.sort()
        inserts.reverse()

        insert_str_ar = []
        incount = 0
        for count,insert in inserts:
            insert_str_ar.append('%s:%s' % (insert,count))
            incount += count

        if big_total > 0:
            ave_mapping = (float(read_ih_acc) / big_total)
        else:
            ave_mapping = 0

        sys.stdout.write('%s\n' % '\t'.join([str(x) for x in (bbc.bam.references[basepos.tid],basepos.pos+1,refbase,basepos.total,ave_mapping,entropy,basepos.a,basepos.c,basepos.g,basepos.t,basepos.n,basepos.deletions,basepos.gaps,incount,','.join(insert_str_ar))]))

    bbc.close()
    if ref:
        ref.close()

if __name__ == '__main__':
    bam = None
    ref = None

    min_qual = 0
    min_count = 0
    mask = 1540
    chrom = None
    start = None
    end = None
    quiet = False

    last = None
    for arg in sys.argv[1:]:
        if last == '-qual':
            min_qual = int(arg)
            last = None
        elif last == '-ref':
            if os.path.exists(arg) and os.path.exists('%s.fai' % arg):
                ref = arg
            else:
                print "Missing FASTA file or index: %s" % arg
                usage()
            last = None
        elif last == '-count':
            min_count = int(arg)
            last = None
        elif last == '-mask':
            mask = int(arg)
            last = None
        elif arg == '-h':
            usage()
        elif arg == '-q':
            quiet = True
        elif arg in ['-qual','-count','-mask','-ref']:
            last = arg
        elif not bam and os.path.exists(arg) and os.path.exists('%s.bai' % arg):
            bam = arg
        elif not ref and os.path.exists(arg) and os.path.exists('%s.fai' % arg):
            ref = arg
        elif not chrom:
            chrom,startend = arg.split(':')
            if '-' in startend:
                start,end = [int(x) for x in startend.split('-')]
            else:
                start = int(startend)
                end = start
            start = start - 1
        else:
            print "Unknown option or missing index: %s" % arg
            usage()

    if not bam:
        usage()
    else:
        bam_basecall(bam,ref,min_qual,min_count,chrom,start,end,mask,quiet)

<commit_msg>Updated docs, added showgaps option (to toggle gaps/introns)<commit_after>#!/usr/bin/env python
"""
Given a BAM file and a genomic reference, for each position covered in the
BAM file, show the reference base, and the number of A/T/C/G's and InDels.

You can also optionally filter out all bases whose quality score is below a
threshold, bases that aren't covered by enough reads, bases that have no
variation compared to reference, or bases whose variation is too low.

The output is a tab-delimited file that contains the following for each base:

chromosome
position (1-based)
reference base
# reads that contain this base
Entropy
# A calls
# C calls
# G calls
# T calls
# deletions
# gaps
# inserts

Entropy is sum(a..t) { p log2 p } where p = freq(+pseudocount) / genomic freq.
pseudo count = genomic freq * sqrt(N)

We use the following genomic frequencies: A 0.3, C 0.2, G 0.2, T 0.3
"""

import os,sys,math,collections
from support.eta import ETA
import pysam

def usage():
    base = os.path.basename(sys.argv[0])
    print __doc__
    print """
Usage: %s {opts} in.bam {chrom:start-end}

Options:
-ref   val    Include reference basecalls from this file
-qual  val    Minimum quality level to use in calculations
              (numeric, Sanger scale) (default 0)

-count val    Report only bases with this minimum number of reads covering it
              (matches, inserts, deletions counted) (default 0)

-mask  val    The bitmask to use for filtering reads (default 1540)

-showgaps     Report gaps/splice-junctions in RNA-seq data

""" % (base)
    sys.exit(1)

__genomic_freq = {'A': 0.3, 'C': 0.2, 'G': 0.2,'T': 0.3}

def calc_entropy(a,c,t,g):
    counts = {'A':a,'C':c,'G':g,'T':t,}

    N = counts['A'] + counts['C'] + counts['G'] + counts['T']
    if N == 0:
        return 0

    N_sqrt = math.sqrt(N)

    count_pseudo = {}
    N_pseudo = 0

    for base in 'ATCG':
        count_pseudo[base] = counts[base] + (__genomic_freq[base] * N_sqrt)
        N_pseudo += count_pseudo[base]

    acc = 0
    for base in 'ATCG':
        p = float(count_pseudo[base]) / N_pseudo / __genomic_freq[base]
        acc += (p * math.log(p,2))

    return acc

MappingRecord = collections.namedtuple('MappingRecord','qpos cigar_op base qual read')
MappingPos = collections.namedtuple('MappingPos','tid pos records')
BasePosition = collections.namedtuple('BasePosition','tid pos total a c g t n deletions gaps insertions reads')

class BamBaseCaller(object):
    def __init__(self, bam_fname, min_qual=0, min_count=0, chrom=None, start=0, end=0, mask=1540,quiet=False):
        self.bam = pysam.Samfile(bam_fname,'rb')
        self.min_qual = min_qual
        self.min_count = 0

        self.chrom = chrom
        self.start = start
        self.end = end

        self.mask = mask
        self.quiet = quiet

        if chrom and start and end:
            def _gen():
                for p in self.bam.fetch(chrom,start,end):
                    yield p
        else:
            def _gen():
                for p in self.bam:
                    yield p
        self._gen = _gen

        self.buffer = None
        self.current_tid = None

    def close(self):
        self.bam.close()

    def _calc_pos(self,tid,pos,records):
        counts = {'A':0,'C':0,'G':0,'T':0,'N':0}
        insertions = {}
        deletions = 0
        gaps = 0
        total = 0
        reads = []

        for record in records:
            qpos,cigar_op,base,qual,read = record
            if cigar_op == 0: # M
                if qual >= self.min_qual and (read.flag & self.mask) == 0 :
                    total += 1
                    reads.append(record)

                    counts[base] += 1
            elif cigar_op == 1: # I
                if qual >= self.min_qual and (read.flag & self.mask) == 0 :
                    reads.append(record)

                    if not base in insertions:
                        insertions[base] = 1
                    else:
                        insertions[base] += 1
            elif cigar_op == 2: # D
                deletions += 1
                #total += 1 # not sure these should be included,
                           # samtools mpileup includes them
                           # IGV doesn't

                reads.append(record)
            elif cigar_op == 3: # N
                gaps += 1
                reads.append(record)

        if total >= self.min_count:
            return BasePosition(tid,pos,total,counts['A'],counts['C'],counts['G'],counts['T'],counts['N'],deletions,gaps,insertions,reads)

    def fetch(self):
        self.current_tid = None
        self.buffer = collections.deque()
        if not self.quiet:
            eta = ETA(0,bamfile=self.bam)
        else:
            eta = None

        for read in self._gen():
            if eta:
                eta.print_status(extra='%s:%s (%s)' % (self.bam.references[read.tid],read.pos,len(self.buffer)),bam_pos=(read.tid,read.pos))

            if self.current_tid != read.tid: # new chromosome
                while self.buffer:
                    tid,pos,records = self.buffer.popleft()
                    yield self._calc_pos(tid,pos,records)

                self.current_tid = read.tid

            # handle all positions that are 5' of the current one
            while self.buffer and read.pos > self.buffer[0].pos:
                tid,pos,records = self.buffer.popleft()
                yield self._calc_pos(tid,pos,records)

            self._push_read(read)

        # flush buffer for the end
        while self.buffer:
            tid,pos,records = self.buffer.popleft()
            yield self._calc_pos(tid,pos,records)

        if eta:
            eta.done()

    def _push_read(self,read):

        if not self.buffer:
            self.buffer.append(MappingPos(read.tid, read.pos, []))

        while self.buffer[-1].pos < read.aend:
            self.buffer.append(MappingPos(read.tid, self.buffer[-1].pos+1, []))

        buf_idx = 0
        while self.buffer[0].pos < read.pos:
            buf_idx += 1

        read_idx = 0
        for op,length in read.cigar:
            if op == 0:
                for i in xrange(length):
                    self.buffer[buf_idx].records.append(MappingRecord(read_idx,op,read.seq[read_idx],read.qual[read_idx],read))
                    buf_idx += 1
                    read_idx += 1
            elif op == 1:
                inseq = ''
                inqual = 0
                for i in xrange(length):
                    inseq += read.seq[read_idx]
                    inqual += ord(read.qual[read_idx])-33
                    read_idx += 1

                inqual = inqual / len(inseq)

                self.buffer[buf_idx].records.append(MappingRecord(read_idx,op,inseq,inqual,read))

            elif op == 2:
                for i in xrange(length):
                    self.buffer[buf_idx].records.append(MappingRecord(read_idx,op,None,None,read))
                    buf_idx += 1

            elif op == 3:
                for i in xrange(length):
                    self.buffer[buf_idx].records.append(MappingRecord(read_idx,op,None,None,read))
                    buf_idx += 1

def bam_basecall(bam_fname,ref_fname,min_qual=0, min_count=0, chrom=None,start=None,end=None,mask=1540,quiet = False, showgaps=False):
    if ref_fname:
        ref = pysam.Fastafile(ref_fname)
    else:
        ref = None

    sys.stdout.write('chrom\tpos\tref\tcount\tave mappings\tentropy\tA\tC\tG\tT\tN\tDeletions\tGaps\tInsertions\tInserts\n')
    bbc = BamBaseCaller(bam_fname,min_qual,min_count,chrom,start,end,mask,quiet)
    for basepos in bbc.fetch():
        if start and end:
            if basepos.pos < start or basepos.pos >= end:
                continue

        big_total = basepos.total + basepos.deletions + len(basepos.insertions)

        if big_total < min_count:
            continue

        if big_total == 0 and not (showgaps and basepos.gaps > 0):
            continue

        if ref:
            refbase = ref.fetch(bbc.bam.references[basepos.tid],basepos.pos,basepos.pos+1).upper()
        else:
            refbase = 'N'

        entropy = calc_entropy(basepos.a,basepos.c,basepos.g,basepos.t)

        read_ih_acc = 0
        for qpos,cigar_op,base,qual,read in basepos.reads:
            if cigar_op in [0,1,2]:
                read_ih_acc += int(read.opt('IH'))

        inserts = []
        for insert in basepos.insertions:
            inserts.append((basepos.insertions[insert],insert))
        inserts.sort()
        inserts.reverse()

        insert_str_ar = []
        incount = 0
        for count,insert in inserts:
            insert_str_ar.append('%s:%s' % (insert,count))
            incount += count

        if big_total > 0:
            ave_mapping = (float(read_ih_acc) / big_total)
        else:
            ave_mapping = 0

        sys.stdout.write('%s\n' % '\t'.join([str(x) for x in (bbc.bam.references[basepos.tid],basepos.pos+1,refbase,basepos.total,ave_mapping,entropy,basepos.a,basepos.c,basepos.g,basepos.t,basepos.n,basepos.deletions,basepos.gaps,incount,','.join(insert_str_ar))]))

    bbc.close()
    if ref:
        ref.close()

if __name__ == '__main__':
    bam = None
    ref = None

    min_qual = 0
    min_count = 0
    mask = 1540
    chrom = None
    start = None
    end = None
    quiet = False
    showgaps = False

    last = None
    for arg in sys.argv[1:]:
        if last == '-qual':
            min_qual = int(arg)
            last = None
        elif last == '-ref':
            if os.path.exists(arg) and os.path.exists('%s.fai' % arg):
                ref = arg
            else:
                print "Missing FASTA file or index: %s" % arg
                usage()
            last = None
        elif last == '-count':
            min_count = int(arg)
            last = None
        elif last == '-mask':
            mask = int(arg)
            last = None
        elif arg == '-h':
            usage()
        elif arg == '-showgaps':
            showgaps = True
        elif arg == '-q':
            quiet = True
        elif arg in ['-qual','-count','-mask','-ref']:
            last = arg
        elif not bam and os.path.exists(arg) and os.path.exists('%s.bai' % arg):
            bam = arg
        elif not ref and os.path.exists(arg) and os.path.exists('%s.fai' % arg):
            ref = arg
        elif not chrom:
            chrom,startend = arg.split(':')
            if '-' in startend:
                start,end = [int(x) for x in startend.split('-')]
            else:
                start = int(startend)
                end = start
            start = start - 1
        else:
            print "Unknown option or missing index: %s" % arg
            usage()

    if not bam:
        usage()
    else:
        bam_basecall(bam,ref,min_qual,min_count,chrom,start,end,mask,quiet,showgaps)

<|endoftext|>


<commit_before>#coding: utf-8
import datetime
from django.utils.datastructures import SortedDict
from calendar import HTMLCalendar, LocaleHTMLCalendar, Calendar
from django import template

register = template.Library()

@register.simple_tag
def gencal(obj_list, year=None, month=None, calendar_class=None):
    """
    Renders a simple calendar of the given month and year if none are
    specified. Accomplishes this by passing the arguments to
    :class:`ListCalendar`

    ::

      {% gencal queryset %}

      {% gencal queryset 1983 12 %}

    :param obj_list: A list of objects to render on the calendar.
    :type obj_list: list.
    :keyword year: Year to render.
    :type year: int.
    :keyword month: Month to render.
    :type month: int.
    :returns: calendar as HTML
    :rtype: str.
    """
    today = datetime.date.today()
    if not year:
        year = today.year
    if not month:
        month = today.month
    if not calendar_class:
        calendar_class = ListCalendar
    return ''.join(calendar_class(obj_list, year=year, month=month).formatmonth(year, month))

class ListCalendar(HTMLCalendar):
    """
    This is a calendar object which accepts a ``list`` argument and a
    ``date_field`` keyword argument.. This class will return an HTML
    calendar with links on the days that are present in the list,
    using ``date_field`` as the lookup.

    It's assumed that the provided list is valid for the given month,
    ie: contains no outside dates.

    Usage::

        >>> from datetime import datetime, timedelta
        >>> from gencal.templatetags.gencal import ListCalendar
        >>> subclass = type('SubListClass', (ListCalendar,),{}) # This is equivalent to a subclass
        >>> subclass.get_link(self, dt) = lambda dt: "/items/%d/%d/%d" % (dt.year, dt.month, dt.day)
        >>> dates = [{'date':datetime.now(), 'name':'test'},
        ... {'date':datetime.now() + timedelta(days=5), 'name':'test2'}]
        >>> lc = subclass(dates)
        >>> print "".join(lc.formatmonth(2009, 01))

January 2009

Mon	Tue	Wed	Thu	Fri	Sat	Sun

29	30	31	1	2	3	4

5	6	7	8	9	10	11

12	13	14	15	16	17	18

19	20	21	22	23	24	25

26	27	28	29	30	31	1

    :param cal_items: A list of items to put in the calendar.
    :type cal_items: list.
    :keyword year: Year to render.
    :type year: int.
    :keyword month: Month to render.
    :type month: int.
    """

    DAY_ABBR = ["пн", "вт", "ср", "чтв", "пт", "сбт", "вс"]
    cssclasses = ["mon", "tue", "wed", "thu", "fri", "sat", "sun"]

    def __init__(self, cal_items, year=None, month=None, *args, **kwargs):
        today = datetime.date.today()
        self.today = today

        if year == None:
            year = today.year
        self.year = year

        if not month:
            month = today.month
        self.month = month

        self.date_field = kwargs.pop('date_field', 'date')

        super(ListCalendar, self).__init__(*args, **kwargs)

        cal_arr = self.monthdatescalendar(year, month)
        month_dict = SortedDict()
        for week in cal_arr:
            for date in week:
                month_dict[date] = []

        for item in cal_items:
            if isinstance(item, dict):
                possible_date = item.get(self.date_field)
            else:
                possible_date = getattr(item, self.date_field)
            if type(possible_date) == datetime.datetime:
                # transform possible_date to a date, not a datetime
                possible_date = possible_date.date()
            if possible_date:
                month_dict[possible_date].append(item)
        self.month_dict = month_dict

    def formatday(self, day, weekday, themonth):
        """
        Return a day as a table cell.

        :arg day: A day to be formatted.
        :type day: date object.
        :arg weekday: Weekday of given day.
        :type weekday: int.
        """
        if day.month != themonth:
            return ""
        if self.month_dict[day]:
            content = '%d' % (self.get_link(day), day.day)
        else:
            content = str(day.day)
        if day == self.today:
            return '%s	' % content
        else:
            return '%s	' % content

    def get_link(self, dt):
        """
        Should return a url to a given date, as represented by ``dt``

        :arg dt: date to turn into a url
        :type dt: date/datetime
        """
        return "#%s" % dt

    def monthdates2calendar(self, year, month):
        """
        Function returns a list containing a list of tuples
        (representing a week), the tuples represent a
        ``datetime.date`` for the given day and the weekday associated
        with it.

        This is something that ought to be implemented in core. We
        have monthdatescalendar, monthdayscalendar,
        monthdays2calendar, but no monthdates2calendar.

        :keyword year: Year to render.
        :type year: int.
        :keyword month: Month to render.
        :type month: int.
        :returns: Tuple of (datetime, weekday).
        :rtype: tuple(datetime, int)
        """
        return [[(dt, dt.weekday()) for dt in week] for week in self.monthdatescalendar(year, month)]

    def formatweekday(self, day):
        """
        Return a weekday name as a table header.
        """
        return '%s	' % (self.cssclasses[day], self.DAY_ABBR[day])

    def formatweekheader(self):
        """
        Return a header for a week as a table row.
        """
        s = ''.join(self.formatweekday(i) for i in self.iterweekdays())
        return '
%s
' % s

    def formatweek(self, theweek, themonth):
        """
        Return a complete week as a table row.
        """
        s = ''.join(self.formatday(d, wd, themonth) for (d, wd) in theweek)
        return '
%s
' % s

    def formatmonth(self, theyear, themonth, withyear=True):
        """
        Return a formatted month as a table.

        Overridden so weeks will use monthdates2calendar, so we have
        access to full date objects, rather than numbers.

        :arg theyear: Year of calendar to render.
        :type theyear: int.
        :arg themonth: Month of calendar to render
        :type themonth: int.
        :keyword withyear: If true, it will show the year in the header.
        :type withyear: bool.
        """
        v = []
        a = v.append

        a('
')
        a('\n')
        a(self.formatmonthname(theyear, themonth, withyear=withyear))
        a('\n')
        a(self.formatweekheader())
        a('\n')
        for week in self.monthdates2calendar(theyear, themonth):
            a(self.formatweek(week, themonth))
            a('\n')
        a('
')
        a('\n')
        return v
<commit_msg>fix gencal<commit_after>#coding: utf-8
import datetime
from django.utils.datastructures import SortedDict
from calendar import HTMLCalendar, LocaleHTMLCalendar, Calendar
from django import template

register = template.Library()

@register.simple_tag
def gencal(obj_list, year=None, month=None, calendar_class=None):
    """
    Renders a simple calendar of the given month and year if none are
    specified. Accomplishes this by passing the arguments to
    :class:`ListCalendar`

    ::

      {% gencal queryset %}

      {% gencal queryset 1983 12 %}

    :param obj_list: A list of objects to render on the calendar.
    :type obj_list: list.
    :keyword year: Year to render.
    :type year: int.
    :keyword month: Month to render.
    :type month: int.
    :returns: calendar as HTML
    :rtype: str.
    """
    today = datetime.date.today()
    if not year:
        year = today.year
    if not month:
        month = today.month
    if not calendar_class:
        calendar_class = ListCalendar
    return ''.join(calendar_class(obj_list, year=year, month=month).formatmonth(year, month))

class ListCalendar(HTMLCalendar):
    """
    This is a calendar object which accepts a ``list`` argument and a
    ``date_field`` keyword argument.. This class will return an HTML
    calendar with links on the days that are present in the list,
    using ``date_field`` as the lookup.

    It's assumed that the provided list is valid for the given month,
    ie: contains no outside dates.

    Usage::

        >>> from datetime import datetime, timedelta
        >>> from gencal.templatetags.gencal import ListCalendar
        >>> subclass = type('SubListClass', (ListCalendar,),{}) # This is equivalent to a subclass
        >>> subclass.get_link(self, dt) = lambda dt: "/items/%d/%d/%d" % (dt.year, dt.month, dt.day)
        >>> dates = [{'date':datetime.now(), 'name':'test'},
        ... {'date':datetime.now() + timedelta(days=5), 'name':'test2'}]
        >>> lc = subclass(dates)
        >>> print "".join(lc.formatmonth(2009, 01))

January 2009

Mon	Tue	Wed	Thu	Fri	Sat	Sun

29	30	31	1	2	3	4

5	6	7	8	9	10	11

12	13	14	15	16	17	18

19	20	21	22	23	24	25

26	27	28	29	30	31	1

    :param cal_items: A list of items to put in the calendar.
    :type cal_items: list.
    :keyword year: Year to render.
    :type year: int.
    :keyword month: Month to render.
    :type month: int.
    """

    DAY_ABBR = ["пн", "вт", "ср", "чтв", "пт", "сбт", "вс"]
    cssclasses = ["mon", "tue", "wed", "thu", "fri", "sat", "sun"]
    MONTH_LIST = ["Январь", "Февраль", "Март", "Апрель", "Май", "Июнь", "Июль",
                  "Август", "Сентябрь", "Октябрь", "Ноябрь", "Декабрь"]

    def __init__(self, cal_items, year=None, month=None, *args, **kwargs):
        today = datetime.date.today()
        self.today = today

        if year == None:
            year = today.year
        self.year = year

        if not month:
            month = today.month
        self.month = month

        self.date_field = kwargs.pop('date_field', 'date')

        super(ListCalendar, self).__init__(*args, **kwargs)

        cal_arr = self.monthdatescalendar(year, month)
        month_dict = SortedDict()
        for week in cal_arr:
            for date in week:
                month_dict[date] = []

        for item in cal_items:
            if isinstance(item, dict):
                possible_date = item.get(self.date_field)
            else:
                possible_date = getattr(item, self.date_field)
            if type(possible_date) == datetime.datetime:
                # transform possible_date to a date, not a datetime
                possible_date = possible_date.date()
            if possible_date:
                month_dict[possible_date].append(item)
        self.month_dict = month_dict

    def formatday(self, day, weekday, themonth):
        """
        Return a day as a table cell.

        :arg day: A day to be formatted.
        :type day: date object.
        :arg weekday: Weekday of given day.
        :type weekday: int.
        """
        if day.month != themonth:
            return ""
        if self.month_dict[day]:
            content = '%d' % (self.get_link(day), day.day)
        else:
            content = str(day.day)
        if day == self.today:
            return '%s	' % content
        else:
            return '%s	' % content

    def get_link(self, dt):
        """
        Should return a url to a given date, as represented by ``dt``

        :arg dt: date to turn into a url
        :type dt: date/datetime
        """
        return "#%s" % dt

    def monthdates2calendar(self, year, month):
        """
        Function returns a list containing a list of tuples
        (representing a week), the tuples represent a
        ``datetime.date`` for the given day and the weekday associated
        with it.

        This is something that ought to be implemented in core. We
        have monthdatescalendar, monthdayscalendar,
        monthdays2calendar, but no monthdates2calendar.

        :keyword year: Year to render.
        :type year: int.
        :keyword month: Month to render.
        :type month: int.
        :returns: Tuple of (datetime, weekday).
        :rtype: tuple(datetime, int)
        """
        return [[(dt, dt.weekday()) for dt in week] for week in self.monthdatescalendar(year, month)]

    def formatweekday(self, day):
        """
        Return a weekday name as a table header.
        """
        return '%s	' % (self.cssclasses[day], self.DAY_ABBR[day])

    def formatweekheader(self):
        """
        Return a header for a week as a table row.
        """
        s = ''.join(self.formatweekday(i) for i in self.iterweekdays())
        return '
%s
' % s

    def formatweek(self, theweek, themonth):
        """
        Return a complete week as a table row.
        """
        s = ''.join(self.formatday(d, wd, themonth) for (d, wd) in theweek)
        return '
%s
' % s

    def formatmonthname(self, year, month, withyear=True):
        return '%s %s	' % (
                self.MONTH_LIST[month - 1], year)

    def formatmonth(self, theyear, themonth, withyear=True):
        """
        Return a formatted month as a table.

        Overridden so weeks will use monthdates2calendar, so we have
        access to full date objects, rather than numbers.

        :arg theyear: Year of calendar to render.
        :type theyear: int.
        :arg themonth: Month of calendar to render
        :type themonth: int.
        :keyword withyear: If true, it will show the year in the header.
        :type withyear: bool.
        """
        v = []
        a = v.append

        a('
')
        a('\n')
        a(self.formatmonthname(theyear, themonth, withyear=withyear))
        a('\n')
        a(self.formatweekheader())
        a('\n')
        for week in self.monthdates2calendar(theyear, themonth):
            a(self.formatweek(week, themonth))
            a('\n')
        a('
')
        a('\n')
        return v
<|endoftext|>


<commit_before>from __future__ import absolute_import
from argparse import ArgumentParser
import re
import inspect
import os
from os.path import join
import json
import imp
import shutil
from .config import conf
from . import experiment
from limix_util.report import BeginEnd
from limix_util.inspect_ import fetch_classes, fetch_functions
import limix_lsf

_workspaces = dict()
def get_workspace(workspace_id):
    if workspace_id not in _workspaces:
        _workspaces[workspace_id] = Workspace(workspace_id)
    return _workspaces[workspace_id]

def get_experiment(workspace_id, experiment_id):
    w = get_workspace(workspace_id)
    if w is None:
        raise Exception('There is no workspace called %s.' % workspace_id)

    e = w.get_experiment(experiment_id)
    if e is None:
        raise Exception('There is no experiment '
                        'called %s in the workspace %s.'
                        % (experiment_id, workspace_id))

    return e

def exists(workspace_id):
    folder = join(conf.get('default', 'base_dir'), workspace_id)
    return os.path.exists(folder)

def get_workspace_ids():
    files = os.listdir(conf.get('default', 'base_dir'))
    return [f for f in files if os.path.isdir(f)]

class Workspace(object):
    def __init__(self, workspace_id):
        super(Workspace, self).__init__()
        self._workspace_id = workspace_id
        self._script_filepath = None
        self._experiments = dict()
        self.force_cache = False

    def rm_experiment(self, experiment_id):
        e = self.get_experiment(experiment_id)
        e.kill_bjobs()
        if os.path.exists(e.folder):
            with BeginEnd('Removing folder'):
                shutil.rmtree(e.folder)

    def get_properties(self):
        try:
            with open(join(self.folder, 'properties.json')) as json_file:
                return json.load(json_file)
        except ValueError as e:
            print(e)
            print("File: %s", join(self.folder, 'properties.json'))
            print('File content:')
            print(open(join(self.folder, 'properties.json')).read())
            os.system("hostname")

    def get_plot_class(self, name):
        return self._get_plot_classes_map()[name]

    def _get_plot_classes_map(self):
        f = join(self.folder, 'plot.json')
        with open(f) as json_file:
            fpaths = json.load(json_file)
            fpaths = [str(fpath) for fpath in fpaths]

        plot_classes_map = dict()
        for fp in fpaths:
            clss = fetch_classes(fp, '.*Plot')
            for clsi in clss:
                plot_classes_map[clsi.__name__] = clsi
        return plot_classes_map

    @property
    def folder(self):
        return join(conf.get('default', 'base_dir'), self._workspace_id)

    def _get_auto_run(self, experiment_id):
        auto_runs_map = self._get_auto_runs_map()
        # if experiment_id not in auto_runs_map:
            # return None
        return auto_runs_map[experiment_id]

    def get_experiment(self, experiment_id):
        if experiment_id not in self._experiments:
            self._setup_experiment(experiment_id)
        return self._experiments[experiment_id]

    def _setup_experiment(self, experiment_id):
        self._experiments[experiment_id] =\
            experiment.Experiment(self._workspace_id, experiment_id,
                                  self.get_properties())
        auto_run = self._get_auto_run(experiment_id)
        if auto_run is None:
            return
        auto_run(self._experiments[experiment_id])
        self._experiments[experiment_id].auto_run_done = True
        self._experiments[experiment_id].finish_setup()

    def _get_auto_runs_map(self):
        fps = self._auto_run_filepaths()
        auto_runs_map = dict()
        for fp in fps:
            funcs = fetch_functions(fp, r'^auto_run_.+$')
            auto_runs_map.update([(f.func_name[9:], f) for f in funcs])
        return auto_runs_map

    @property
    def dataset_folder(self):
        return os.path.join(self.folder, 'dataset')

    def _auto_run_filepaths(self):
        f = join(self.folder, 'auto_run.json')
        if not os.path.exists(f):
            return []
        with open(f) as json_file:
            fpaths = json.load(json_file)
            fpaths = [str(fpath) for fpath in fpaths]
            return fpaths

    def auto_run(self, args):
        filepaths = self._auto_run_filepaths()
        for fp in filepaths:
            self._auto_run(fp, args)

    def _auto_run(self, script_filepath, args):

        p = ArgumentParser()
        p.add_argument('--dryrun', dest='dryrun', action='store_true')
        p.add_argument('--no-dryrun', dest='dryrun', action='store_false')
        p.set_defaults(dryrun=False, parallel=True)

        args = p.parse_args(args)

        lista = self._get_generate_tasks(script_filepath)
        for (experiment_id, func) in lista:
            exp = self.get_experiment(experiment_id)
            func(exp)

            try:
                exp.submit_jobs(args.dryrun)
            except KeyboardInterrupt:
                shutil.rmtree(exp.folder)
                raise

            if args.dryrun:
                shutil.rmtree(exp.folder)

    def remove(self, experiment_id, jobs_too=False):
        if experiment_id is None:
            print('Nothing to remove.')
            return
        if len(experiment_id) == 0:
            print('Nothing to remove.')
            return
        with BeginEnd('Removing folder'):
            shutil.rmtree(join(self.folder, experiment_id))

        if jobs_too:
            bgroup = '/%s/%s' % (self._workspace_id, experiment_id)
            if exists(bgroup):
                limix_lsf.util.kill_group(bgroup, True)

    def _call_job(self, args):
        experiment_id = args.experiment_id
        (experiment_id, auto_run_func) = self._get_generate_tasks(experiment_id)
        exp = self.get_experiment(experiment_id)
        exp.script_filepath = self._script_filepath
        auto_run_func(exp)
        exp.start(check_existence=False)

    def _get_generate_tasks(self, script_filepath, experiment_id=None):
        script_name = os.path.basename(script_filepath)
        mod_name = os.path.splitext(script_name)[0]
        mod = imp.load_source(mod_name, script_filepath)

        lista = []
        for func in inspect.getmembers(mod, inspect.isfunction):
            func_name = func[0]
            m = re.match(r'^auto_run_(.+)$', func_name)
            if m:
                exp_name = m.group(1)
                if experiment_id is None:
                    lista.append((exp_name, func[1]))
                elif exp_name == experiment_id:
                    return exp_name, func[1]

        return lista

    def __str__(self):
        from tabulate import tabulate

        files = os.listdir(self.folder)
        exp_ids = [f for f in files if (os.path.isdir(join(self.folder, f)) and
                                        join(self.folder, f,
                                             'generate_tasks.txt'))]

        auto_run_err = set()
        tasks_setup_err = set()
        finish_setup_err = set()
        imp_err = set()
        suc = set()

        for ei in exp_ids:
            try:
                e = self.get_experiment(ei)
                if not e.auto_run_done:
                    auto_run_err.add(ei)
                elif not e.tasks_setup_done:
                    tasks_setup_err.add(ei)
                elif not e.finish_setup_done:
                    finish_setup_err.add(ei)
                else:
                    suc.add(ei)
            except ImportError:
                imp_err.add(ei)

        table = []
        auto_run_fps = '\n'.join(self._auto_run_filepaths())
        table.append(['auto_run files', auto_run_fps])
        table.append(['# experiments', str(len(exp_ids))])

        table.append(["# auto_run err", len(auto_run_err)])
        table.append(["# tasks setup err", len(tasks_setup_err)])
        table.append(["# finish setup err", len(finish_setup_err)])
        table.append(["# import err", len(imp_err)])
        table.append(["# successful", len(suc)])

        msg = tabulate(table)

        return msg

def _get_auto_runs_map(script_filepath):
    script_name = os.path.basename(script_filepath)
    mod_name = os.path.splitext(script_name)[0]
    mod = imp.load_source(mod_name, script_filepath)

    auto_run_map = dict()
    for func in inspect.getmembers(mod, inspect.isfunction):
        func_name = func[0]
        m = re.match(r'^auto_run_(.+)$', func_name)
        if m:
            exp_name = m.group(1)
            auto_run_map.update({exp_name:func[1]})

    return auto_run_map
<commit_msg>queue support<commit_after>from __future__ import absolute_import
from argparse import ArgumentParser
import re
import inspect
import os
from os.path import join
import json
import imp
import shutil
from .config import conf
from . import experiment
from limix_util.report import BeginEnd
from limix_util.inspect_ import fetch_classes, fetch_functions
import limix_lsf

_workspaces = dict()
def get_workspace(workspace_id):
    if workspace_id not in _workspaces:
        _workspaces[workspace_id] = Workspace(workspace_id)
    return _workspaces[workspace_id]

def get_experiment(workspace_id, experiment_id):
    w = get_workspace(workspace_id)
    if w is None:
        raise Exception('There is no workspace called %s.' % workspace_id)

    e = w.get_experiment(experiment_id)
    if e is None:
        raise Exception('There is no experiment '
                        'called %s in the workspace %s.'
                        % (experiment_id, workspace_id))

    return e

def exists(workspace_id):
    folder = join(conf.get('default', 'base_dir'), workspace_id)
    return os.path.exists(folder)

def get_workspace_ids():
    files = os.listdir(conf.get('default', 'base_dir'))
    return [f for f in files if os.path.isdir(f)]

class Workspace(object):
    def __init__(self, workspace_id):
        super(Workspace, self).__init__()
        self._workspace_id = workspace_id
        self._script_filepath = None
        self._experiments = dict()
        self.force_cache = False

    def rm_experiment(self, experiment_id):
        e = self.get_experiment(experiment_id)
        e.kill_bjobs()
        if os.path.exists(e.folder):
            with BeginEnd('Removing folder'):
                shutil.rmtree(e.folder)

    def get_properties(self):
        try:
            with open(join(self.folder, 'properties.json')) as json_file:
                return json.load(json_file)
        except ValueError as e:
            print(e)
            print("File: %s", join(self.folder, 'properties.json'))
            print('File content:')
            print(open(join(self.folder, 'properties.json')).read())
            os.system("hostname")

    def get_plot_class(self, name):
        return self._get_plot_classes_map()[name]

    def _get_plot_classes_map(self):
        f = join(self.folder, 'plot.json')
        with open(f) as json_file:
            fpaths = json.load(json_file)
            fpaths = [str(fpath) for fpath in fpaths]

        plot_classes_map = dict()
        for fp in fpaths:
            clss = fetch_classes(fp, '.*Plot')
            for clsi in clss:
                plot_classes_map[clsi.__name__] = clsi
        return plot_classes_map

    @property
    def folder(self):
        return join(conf.get('default', 'base_dir'), self._workspace_id)

    def _get_auto_run(self, experiment_id):
        auto_runs_map = self._get_auto_runs_map()
        # if experiment_id not in auto_runs_map:
            # return None
        return auto_runs_map[experiment_id]

    def get_experiment(self, experiment_id):
        if experiment_id not in self._experiments:
            self._setup_experiment(experiment_id)
        return self._experiments[experiment_id]

    def _setup_experiment(self, experiment_id):
        self._experiments[experiment_id] =\
            experiment.Experiment(self._workspace_id, experiment_id,
                                  self.get_properties())
        auto_run = self._get_auto_run(experiment_id)
        if auto_run is None:
            return
        auto_run(self._experiments[experiment_id])
        self._experiments[experiment_id].auto_run_done = True
        self._experiments[experiment_id].finish_setup()

    def _get_auto_runs_map(self):
        fps = self._auto_run_filepaths()
        auto_runs_map = dict()
        for fp in fps:
            funcs = fetch_functions(fp, r'^auto_run_.+$')
            auto_runs_map.update([(f.func_name[9:], f) for f in funcs])
        return auto_runs_map

    @property
    def dataset_folder(self):
        return os.path.join(self.folder, 'dataset')

    def _auto_run_filepaths(self):
        f = join(self.folder, 'auto_run.json')
        if not os.path.exists(f):
            return []
        with open(f) as json_file:
            fpaths = json.load(json_file)
            fpaths = [str(fpath) for fpath in fpaths]
            return fpaths

    def auto_run(self, args):
        filepaths = self._auto_run_filepaths()
        for fp in filepaths:
            self._auto_run(fp, args)

    def _auto_run(self, script_filepath, args):

        p = ArgumentParser()
        p.add_argument('--dryrun', dest='dryrun', action='store_true')
        p.add_argument('--no-dryrun', dest='dryrun', action='store_false')
        p.add_argument('--queue', default=None)
        p.set_defaults(dryrun=False, parallel=True)

        args = p.parse_args(args)

        lista = self._get_generate_tasks(script_filepath)
        for (experiment_id, func) in lista:
            exp = self.get_experiment(experiment_id)
            func(exp)

            try:
                exp.submit_jobs(args.dryrun, queue=args.queue)
            except KeyboardInterrupt:
                shutil.rmtree(exp.folder)
                raise

            if args.dryrun:
                shutil.rmtree(exp.folder)

    def remove(self, experiment_id, jobs_too=False):
        if experiment_id is None:
            print('Nothing to remove.')
            return
        if len(experiment_id) == 0:
            print('Nothing to remove.')
            return
        with BeginEnd('Removing folder'):
            shutil.rmtree(join(self.folder, experiment_id))

        if jobs_too:
            bgroup = '/%s/%s' % (self._workspace_id, experiment_id)
            if exists(bgroup):
                limix_lsf.util.kill_group(bgroup, True)

    def _call_job(self, args):
        experiment_id = args.experiment_id
        (experiment_id, auto_run_func) = self._get_generate_tasks(experiment_id)
        exp = self.get_experiment(experiment_id)
        exp.script_filepath = self._script_filepath
        auto_run_func(exp)
        exp.start(check_existence=False)

    def _get_generate_tasks(self, script_filepath, experiment_id=None):
        script_name = os.path.basename(script_filepath)
        mod_name = os.path.splitext(script_name)[0]
        mod = imp.load_source(mod_name, script_filepath)

        lista = []
        for func in inspect.getmembers(mod, inspect.isfunction):
            func_name = func[0]
            m = re.match(r'^auto_run_(.+)$', func_name)
            if m:
                exp_name = m.group(1)
                if experiment_id is None:
                    lista.append((exp_name, func[1]))
                elif exp_name == experiment_id:
                    return exp_name, func[1]

        return lista

    def __str__(self):
        from tabulate import tabulate

        files = os.listdir(self.folder)
        exp_ids = [f for f in files if (os.path.isdir(join(self.folder, f)) and
                                        join(self.folder, f,
                                             'generate_tasks.txt'))]

        auto_run_err = set()
        tasks_setup_err = set()
        finish_setup_err = set()
        imp_err = set()
        suc = set()

        for ei in exp_ids:
            try:
                e = self.get_experiment(ei)
                if not e.auto_run_done:
                    auto_run_err.add(ei)
                elif not e.tasks_setup_done:
                    tasks_setup_err.add(ei)
                elif not e.finish_setup_done:
                    finish_setup_err.add(ei)
                else:
                    suc.add(ei)
            except ImportError:
                imp_err.add(ei)

        table = []
        auto_run_fps = '\n'.join(self._auto_run_filepaths())
        table.append(['auto_run files', auto_run_fps])
        table.append(['# experiments', str(len(exp_ids))])

        table.append(["# auto_run err", len(auto_run_err)])
        table.append(["# tasks setup err", len(tasks_setup_err)])
        table.append(["# finish setup err", len(finish_setup_err)])
        table.append(["# import err", len(imp_err)])
        table.append(["# successful", len(suc)])

        msg = tabulate(table)

        return msg

def _get_auto_runs_map(script_filepath):
    script_name = os.path.basename(script_filepath)
    mod_name = os.path.splitext(script_name)[0]
    mod = imp.load_source(mod_name, script_filepath)

    auto_run_map = dict()
    for func in inspect.getmembers(mod, inspect.isfunction):
        func_name = func[0]
        m = re.match(r'^auto_run_(.+)$', func_name)
        if m:
            exp_name = m.group(1)
            auto_run_map.update({exp_name:func[1]})

    return auto_run_map
<|endoftext|>


<commit_before>from django.contrib.contenttypes.models import ContentType
from django.db import models, transaction
from manager_utils import ManagerUtilsManager, sync

from data_schema.convert_value import convert_value
from data_schema.field_schema_type import FieldSchemaType

class DataSchemaManager(ManagerUtilsManager):
    """
    A model manager for data schemas. Caches related attributes of data schemas.
    """
    def get_queryset(self):
        return super(DataSchemaManager, self).get_queryset().select_related(
            'model_content_type').prefetch_related('fieldschema_set')

class DataSchema(models.Model):
    """Define a schema information about a unit of data, such as a
    dictionary, list or model.

    ``FieldSchema`` objects are created referencing a ``DataSchema``
    to describe individual fields within the unit of data.
    """
    # The content type of the django model for which this schema is related. If None, this schema is
    # for a dictionary of data.
    model_content_type = models.ForeignKey(ContentType, null=True, default=None)

    # A custom model manager that caches objects
    objects = DataSchemaManager()

    @transaction.atomic
    def update(self, **updates):
        """
        Takes a template of updates and updates the data schema. The following are possible keyword
        arguments.

        'model_content_type': The string path of a model that this schema represents (or None),
        'fieldschema_set': [{
            'field_key': The field key of the field,
            'display_name': The display name for the field (or the 'field_key' by default),
            'field_type': The FieldSchemaType type of the field,
            'uniqueness_order': The order this field is in the uniquessness constraint (or None by default),
            'field_position': The position of this field if it can be parsed by an array (or None by default),
            'field_format': The format of this field (or None by default),
            'default_value': The default value of this field (or None by default),
            'fieldoption_set': The set of options for the field schema (optional),
        }, {
            Additional field schemas...
        }]
        """
        if 'model_content_type' in updates:
            self.model_content_type = updates['model_content_type']

        self.save()

        if 'fieldschema_set' in updates:
            # Sync the field schema models
            sync(self.fieldschema_set.all(), [
                FieldSchema(
                    data_schema=self,
                    field_key=fs_values['field_key'],
                    display_name=fs_values.get('display_name', ''),
                    field_type=fs_values['field_type'],
                    uniqueness_order=fs_values.get('uniqueness_order', None),
                    field_position=fs_values.get('field_position', None),
                    field_format=fs_values.get('field_format', None),
                    default_value=fs_values.get('default_value', None),
                    has_options='fieldoption_set' in fs_values and fs_values['fieldoption_set'],
                    transform_case=fs_values.get('transform_case', None),
                )
                for fs_values in updates['fieldschema_set']
            ], ['field_key'], [
                'display_name', 'field_key', 'field_type', 'uniqueness_order', 'field_position',
                'field_format', 'default_value', 'has_options', 'case'
            ])

            # Sync the options of the field schema models if they are present
            for fs_values in updates['fieldschema_set']:
                if 'fieldoption_set' in fs_values:
                    fs = self.fieldschema_set.get(field_key=fs_values['field_key'])
                    sync(fs.fieldoption_set.all(), [
                        FieldOption(
                            field_schema=fs,
                            value=f_option,
                        )
                        for f_option in fs_values['fieldoption_set']
                    ], ['value'], ['value'])

    def get_unique_fields(self):
        """
        Gets all of the fields that create the uniqueness constraint for a metric record.
        """
        if not hasattr(self, '_unique_fields'):
            # Instead of querying the reverse relationship directly, assume that it has been cached
            # with prefetch_related and go through all fields.
            setattr(self, '_unique_fields', [
                field for field in self.fieldschema_set.all() if field.uniqueness_order is not None
            ])
            self._unique_fields.sort(key=lambda k: k.uniqueness_order or 0)
        return self._unique_fields

    def get_fields(self):
        """
        Gets all fields in the schema. Note - dont use django's order_by since we are caching the fieldschema_set
        beforehand.
        """
        return sorted(self.fieldschema_set.all(), key=lambda k: k.field_position or 0)

    def _get_field_map(self):
        """
        Returns a cached mapping of field keys to their field schemas.
        """
        if not hasattr(self, '_field_map'):
            self._field_map = {
                field.field_key: field for field in self.get_fields()
            }
        return self._field_map

    def get_value(self, obj, field_key):
        """
        Given an object and a field key, return the value of the field in the object.
        """
        try:
            return self._get_field_map()[field_key].get_value(obj)
        except Exception as e:
            # Attach additional information to the exception to make higher level error handling easier
            e.field_key = field_key
            raise e

    def set_value(self, obj, field_key, value):
        """
        Given an object and a field key, set the value of the field in the object.
        """
        return self._get_field_map()[field_key].set_value(obj, value)

class FieldSchema(models.Model):
    """
    Specifies the schema for a field in a piece of data.
    """
    class Meta:
        unique_together = ('data_schema', 'field_key')

    # The data schema to which this field belongs
    data_schema = models.ForeignKey(DataSchema)

    # The key for the field in the data
    field_key = models.CharField(max_length=64)

    # Optional way to display the field. defaults to the field_key
    display_name = models.CharField(max_length=64, blank=True, default='')

    # The order in which this field appears in the UID for the record. It is null if it does
    # not appear in the uniqueness constraint
    uniqueness_order = models.IntegerField(null=True)

    # The position of the field. This ordering is relevant when parsing a list of fields into
    # a dictionary with the field names as keys
    field_position = models.IntegerField(null=True)

    # The type of field. The available choices are present in the FieldSchemaType class
    field_type = models.CharField(max_length=32, choices=FieldSchemaType.choices())

    # If the field is a string and needs to be converted to another type, this string specifies
    # the format for a field
    field_format = models.CharField(null=True, blank=True, default=None, max_length=64)

    # This field provides a default value to be used for the field in the case that it is None.
    default_value = models.CharField(null=True, blank=True, default=None, max_length=128)

    # Flag to indicate if setting values should be validated against an option list
    has_options = models.BooleanField(default=False)

    # Only applies to string data types. Valid options are upper and lower
    transform_case = models.CharField(null=True, default=None, blank=True, max_length=5)

    # Use django manager utils to manage FieldSchema objects
    objects = ManagerUtilsManager()

    def set_value(self, obj, value):
        """
        Given an object, set the value of the field in that object.
        """
        if self.has_options:
            # Build a set of possible values by calling self.get_value on the stored value options. This will
            # make sure they are the right data type because they are stored as strings
            values = set(
                self.get_value({
                    self.field_key: field_option.value
                })
                for field_option in self.fieldoption_set.all()
            )
            if value not in values:
                raise Exception('Invalid option for {0}'.format(self.field_key))

        if isinstance(obj, list):
            obj[self.field_position] = value
        elif isinstance(obj, dict):
            obj[self.field_key] = value
        else:
            setattr(obj, self.field_key, value)

    def get_value(self, obj):
        """
        Given an object, return the value of the field in that object.
        """
        if isinstance(obj, list):
            value = obj[self.field_position] if 0 <= self.field_position < len(obj) else None
        elif isinstance(obj, dict):
            value = obj[self.field_key] if self.field_key in obj else None
        else:
            value = getattr(obj, self.field_key) if hasattr(obj, self.field_key) else None

        return convert_value(self.field_type, value, self.field_format, self.default_value, self.case)

    def save(self, *args, **kwargs):
        if not self.display_name:
            self.display_name = self.field_key
        super(FieldSchema, self).save(*args, **kwargs)

class FieldOption(models.Model):
    """
    Specifies a set of possible values that a field schema can have
    """
    field_schema = models.ForeignKey(FieldSchema)
    value = models.CharField(max_length=1024)

    class Meta:
        unique_together = ('field_schema', 'value')
<commit_msg>property name fix<commit_after>from django.contrib.contenttypes.models import ContentType
from django.db import models, transaction
from manager_utils import ManagerUtilsManager, sync

from data_schema.convert_value import convert_value
from data_schema.field_schema_type import FieldSchemaType

class DataSchemaManager(ManagerUtilsManager):
    """
    A model manager for data schemas. Caches related attributes of data schemas.
    """
    def get_queryset(self):
        return super(DataSchemaManager, self).get_queryset().select_related(
            'model_content_type').prefetch_related('fieldschema_set')

class DataSchema(models.Model):
    """Define a schema information about a unit of data, such as a
    dictionary, list or model.

    ``FieldSchema`` objects are created referencing a ``DataSchema``
    to describe individual fields within the unit of data.
    """
    # The content type of the django model for which this schema is related. If None, this schema is
    # for a dictionary of data.
    model_content_type = models.ForeignKey(ContentType, null=True, default=None)

    # A custom model manager that caches objects
    objects = DataSchemaManager()

    @transaction.atomic
    def update(self, **updates):
        """
        Takes a template of updates and updates the data schema. The following are possible keyword
        arguments.

        'model_content_type': The string path of a model that this schema represents (or None),
        'fieldschema_set': [{
            'field_key': The field key of the field,
            'display_name': The display name for the field (or the 'field_key' by default),
            'field_type': The FieldSchemaType type of the field,
            'uniqueness_order': The order this field is in the uniquessness constraint (or None by default),
            'field_position': The position of this field if it can be parsed by an array (or None by default),
            'field_format': The format of this field (or None by default),
            'default_value': The default value of this field (or None by default),
            'fieldoption_set': The set of options for the field schema (optional),
        }, {
            Additional field schemas...
        }]
        """
        if 'model_content_type' in updates:
            self.model_content_type = updates['model_content_type']

        self.save()

        if 'fieldschema_set' in updates:
            # Sync the field schema models
            sync(self.fieldschema_set.all(), [
                FieldSchema(
                    data_schema=self,
                    field_key=fs_values['field_key'],
                    display_name=fs_values.get('display_name', ''),
                    field_type=fs_values['field_type'],
                    uniqueness_order=fs_values.get('uniqueness_order', None),
                    field_position=fs_values.get('field_position', None),
                    field_format=fs_values.get('field_format', None),
                    default_value=fs_values.get('default_value', None),
                    has_options='fieldoption_set' in fs_values and fs_values['fieldoption_set'],
                    transform_case=fs_values.get('transform_case', None),
                )
                for fs_values in updates['fieldschema_set']
            ], ['field_key'], [
                'display_name', 'field_key', 'field_type', 'uniqueness_order', 'field_position',
                'field_format', 'default_value', 'has_options', 'case'
            ])

            # Sync the options of the field schema models if they are present
            for fs_values in updates['fieldschema_set']:
                if 'fieldoption_set' in fs_values:
                    fs = self.fieldschema_set.get(field_key=fs_values['field_key'])
                    sync(fs.fieldoption_set.all(), [
                        FieldOption(
                            field_schema=fs,
                            value=f_option,
                        )
                        for f_option in fs_values['fieldoption_set']
                    ], ['value'], ['value'])

    def get_unique_fields(self):
        """
        Gets all of the fields that create the uniqueness constraint for a metric record.
        """
        if not hasattr(self, '_unique_fields'):
            # Instead of querying the reverse relationship directly, assume that it has been cached
            # with prefetch_related and go through all fields.
            setattr(self, '_unique_fields', [
                field for field in self.fieldschema_set.all() if field.uniqueness_order is not None
            ])
            self._unique_fields.sort(key=lambda k: k.uniqueness_order or 0)
        return self._unique_fields

    def get_fields(self):
        """
        Gets all fields in the schema. Note - dont use django's order_by since we are caching the fieldschema_set
        beforehand.
        """
        return sorted(self.fieldschema_set.all(), key=lambda k: k.field_position or 0)

    def _get_field_map(self):
        """
        Returns a cached mapping of field keys to their field schemas.
        """
        if not hasattr(self, '_field_map'):
            self._field_map = {
                field.field_key: field for field in self.get_fields()
            }
        return self._field_map

    def get_value(self, obj, field_key):
        """
        Given an object and a field key, return the value of the field in the object.
        """
        try:
            return self._get_field_map()[field_key].get_value(obj)
        except Exception as e:
            # Attach additional information to the exception to make higher level error handling easier
            e.field_key = field_key
            raise e

    def set_value(self, obj, field_key, value):
        """
        Given an object and a field key, set the value of the field in the object.
        """
        return self._get_field_map()[field_key].set_value(obj, value)

class FieldSchema(models.Model):
    """
    Specifies the schema for a field in a piece of data.
    """
    class Meta:
        unique_together = ('data_schema', 'field_key')

    # The data schema to which this field belongs
    data_schema = models.ForeignKey(DataSchema)

    # The key for the field in the data
    field_key = models.CharField(max_length=64)

    # Optional way to display the field. defaults to the field_key
    display_name = models.CharField(max_length=64, blank=True, default='')

    # The order in which this field appears in the UID for the record. It is null if it does
    # not appear in the uniqueness constraint
    uniqueness_order = models.IntegerField(null=True)

    # The position of the field. This ordering is relevant when parsing a list of fields into
    # a dictionary with the field names as keys
    field_position = models.IntegerField(null=True)

    # The type of field. The available choices are present in the FieldSchemaType class
    field_type = models.CharField(max_length=32, choices=FieldSchemaType.choices())

    # If the field is a string and needs to be converted to another type, this string specifies
    # the format for a field
    field_format = models.CharField(null=True, blank=True, default=None, max_length=64)

    # This field provides a default value to be used for the field in the case that it is None.
    default_value = models.CharField(null=True, blank=True, default=None, max_length=128)

    # Flag to indicate if setting values should be validated against an option list
    has_options = models.BooleanField(default=False)

    # Only applies to string data types. Valid options are upper and lower
    transform_case = models.CharField(null=True, default=None, blank=True, max_length=5)

    # Use django manager utils to manage FieldSchema objects
    objects = ManagerUtilsManager()

    def set_value(self, obj, value):
        """
        Given an object, set the value of the field in that object.
        """
        if self.has_options:
            # Build a set of possible values by calling self.get_value on the stored value options. This will
            # make sure they are the right data type because they are stored as strings
            values = set(
                self.get_value({
                    self.field_key: field_option.value
                })
                for field_option in self.fieldoption_set.all()
            )
            if value not in values:
                raise Exception('Invalid option for {0}'.format(self.field_key))

        if isinstance(obj, list):
            obj[self.field_position] = value
        elif isinstance(obj, dict):
            obj[self.field_key] = value
        else:
            setattr(obj, self.field_key, value)

    def get_value(self, obj):
        """
        Given an object, return the value of the field in that object.
        """
        if isinstance(obj, list):
            value = obj[self.field_position] if 0 <= self.field_position < len(obj) else None
        elif isinstance(obj, dict):
            value = obj[self.field_key] if self.field_key in obj else None
        else:
            value = getattr(obj, self.field_key) if hasattr(obj, self.field_key) else None

        return convert_value(self.field_type, value, self.field_format, self.default_value, self.transform_case)

    def save(self, *args, **kwargs):
        if not self.display_name:
            self.display_name = self.field_key
        super(FieldSchema, self).save(*args, **kwargs)

class FieldOption(models.Model):
    """
    Specifies a set of possible values that a field schema can have
    """
    field_schema = models.ForeignKey(FieldSchema)
    value = models.CharField(max_length=1024)

    class Meta:
        unique_together = ('field_schema', 'value')
<|endoftext|>