<commit_before>#!/usr/bin/env python

import argparse
import re
import subprocess
import sys

useCSV = False

def main(arguments):
    parser = argparse.ArgumentParser(
        description='Analyze the code size in a binary')
    parser.add_argument('-arch', type=str,
                        help='the arch to look at', default='arm64')
    parser.add_argument('-categorize', action='store_true',
                        help='categorize symbols', dest='build_categories',
                        default=False)
    parser.add_argument('-list-category', type=str,
                        help='list symbols in category')
    parser.add_argument('-csv', dest='use_csv', action='store_true',
                        help='print results as csv')
    parser.add_argument('-uncategorized', action='store_true',
                        help='show all uncategorized symbols',
                        dest='show_uncategorized',
                        default=False)
    parser.add_argument('bin', help='the binary')
    parser.set_defaults(use_csv=False)

    args = parser.parse_args(arguments)
    if args.use_csv:
        global useCSV
        useCSV = True
        print("Using csv")

    segments = parse_segments(args.bin, args.arch)

    if args.build_categories:
        categorize(segments)
    elif args.show_uncategorized:
        uncategorized(segments)
    elif args.list_category:
        list_category(segments, args.list_category)
    else:
        show_all(segments)

class Symbol(object):
    def __init__(self, name, mangled_name, size):
        self.name = name
        self.mangled_name = mangled_name
        self.count = 1
        self.size = int(size)

def get_symbol_size(sym):
    return sym.size

class Segment(object):
    def __init__(self, name):
        self.name = name
        self.sections = []

class Section(object):
    def __init__(self, name, size):
        self.name = name
        self.size = size
        self.symbols = []

class Category(object):
    def __init__(self, name):
        self.name = name
        self.size = 0
        self.symbols = []

    def add(self, symbol):
        self.symbols.append(symbol)
        self.size += symbol.size

class Categories(object):
    def __init__(self):
        self.category_matching = [
            ['Objective-C function', re.compile(r'.*[+-]\[')],
            ['C++', re.compile(r'_+swift')],
            ['Merged function', re.compile(r'merged ')],
            ['Key path', re.compile(r'key path')],
            ['Function signature specialization',
                re.compile(r'function signature specialization')],
            ['Generic specialization', re.compile(r'generic specialization')],
            ['Reabstraction thunk helper',
                re.compile(r'reabstraction thunk helper')],
            ['vtable thunk', re.compile(r'vtable thunk for')],
            ['@objc thunk', re.compile(r'@objc')],
            ['@nonobjc thunk', re.compile(r'@nonobjc')],
            ['Value witness', re.compile(r'.*value witness for')],
            ['Block copy helper', re.compile(r'_block_copy_helper')],
            ['Block destroy helper', re.compile(r'_block_destroy_helper')],
            ['Block literal global', re.compile(r'___block_literal_global')],
            ['Destroy helper block', re.compile(r'___destroy_helper_block')],
            ['Copy helper block', re.compile(r'___copy_helper_block')],
            ['Object destroy', re.compile(r'_objectdestroy')],
            ['Partial apply forwarder',
                re.compile(r'partial apply forwarder')],
            ['Closure function', re.compile(r'closure #')],
            ['ObjC metadata update function',
                re.compile(r'ObjC metadata update function for')],
            ['Variable initialization expression',
                re.compile(r'variable initialization expression of')],
            ['Global initialization', re.compile(r'_globalinit_')],
            ['Unnamed', re.compile(r'___unnamed_')],
            ['Dyld stubs', re.compile(r'DYLD-STUB\$')],
            ['Witness table accessor',
                re.compile(r'.*witness table accessor for')],
            ['Protocol witness', re.compile(r'protocol witness for')],
            ['Outlined variable', re.compile(r'outlined variable #')],
            ['Outlined value function (copy,destroy,release...)',
                re.compile(r'outlined')],
            ['_symbolic', re.compile(r'_symbolic')],
            ['_associated conformance',
                re.compile(r'_associated conformance')],
            ['Direct field offset', re.compile(r'direct field offset for')],
            ['Value witness tables', re.compile(r'.*value witness table')],
            ['Protocol witness table',
                re.compile(r'.*protocol witness table for')],
            ['Protocol conformance descriptor',
                re.compile(r'protocol conformance descriptor for')],
            ['Lazy protocol witness table cache var',
                re.compile(
                    r'lazy protocol witness table cache variable for type')],
            ['Nominal type descriptor',
                re.compile(r'nominal type descriptor for')],
            ['ObjC class', re.compile(r'_OBJC_CLASS_')],
            ['ObjC metaclass', re.compile(r'_OBJC_METACLASS')],
            ['ObjC ivar', re.compile(r'_OBJC_IVAR')],
            ['Metaclass', re.compile(r'metaclass for')],
            ['Block descriptor', re.compile(r'_+block_descriptor')],
            ['Extension descriptor', re.compile(r'extension descriptor')],
            ['Module descriptor', re.compile(r'module descriptor')],
            ['Associated type descriptor',
                re.compile(r'associated type descriptor for')],
            ['Associated conformance descriptor',
                re.compile(r'associated conformance descriptor for')],
            ['Protocol descriptor', re.compile(r'protocol descriptor for')],
            ['Base conformance descriptor',
                re.compile(r'base conformance descriptor for')],
            ['Protocol requirements base descriptor',
                re.compile(r'protocol requirements base descriptor for')],
            ['Property descriptor', re.compile(r'property descriptor for')],
            ['Method descriptor', re.compile(r'method descriptor for')],
            ['Anonymous descriptor', re.compile(r'anonymous descriptor')],
            ['Type metadata accessor',
                re.compile(r'.*type metadata accessor')],
            ['Type metadata', re.compile(r'.*type metadata')],
            ['Reflection metadata descriptor',
                re.compile(r'reflection metadata .* descriptor')],
        ]

        self.category_mangled_matching = [
            ['Swift variable storage', re.compile(r'^_\$s.*[v][p][Z]?$')],
            ['Swift constructor', re.compile(r'^_\$s.*[f][cC]$')],
            ['Swift initializer', re.compile(r'^_\$s.*[f][ie]$')],
            ['Swift destructor/destroyer', re.compile(r'^_\$s.*[f][dDE]$')],
            ['Swift getter', re.compile(r'^_\$s.*[iv][gG]$')],
            ['Swift setter', re.compile(r'^_\$s.*[iv][swW]$')],
            ['Swift materializeForSet', re.compile(r'^_\$s.*[iv][m]$')],
            ['Swift modify', re.compile(r'^_\$s.*[iv][M]$')],
            ['Swift read', re.compile(r'^_\$s.*[iv][r]$')],
            ['Swift addressor', re.compile(r'^_\$s.*[iv][al][uOop]$')],
            ['Swift function', re.compile(r'^_\$s.*F$')],
            ['Swift unknown', re.compile(r'^_\$s.*')],
        ]
        self.categories = {}

    def categorize_by_name(self, symbol):
        for c in self.category_matching:
            if c[1].match(symbol.name):
                return c[0]
        return None

    def categorize_by_mangled_name(self, symbol):
        for c in self.category_mangled_matching:
            if c[1].match(symbol.mangled_name):
                return c[0]
        return None

    def add_symbol(self, category_name, symbol):
        existing_category = self.categories.get(category_name)
        if existing_category:
            existing_category.add(symbol)
        else:
            new_category = Category(category_name)
            new_category.add(symbol)
            self.categories[category_name] = new_category

    def add(self, symbol):
        category_name = self.categorize_by_name(symbol)
        if category_name:
            self.add_symbol(category_name, symbol)
            return
        category_name = self.categorize_by_mangled_name(symbol)
        if category_name:
            self.add_symbol(category_name, symbol)
        else:
            self.add_symbol('Unknown', symbol)

    def categorize(self, symbols):
        for sym in symbols:
            self.add(sym)

    def print_summary(self, section_size):
        names = [c[0] for c in self.category_matching]
        names.extend([c[0] for c in self.category_mangled_matching])
        names.append('Unknown')
        total_size = 0
        sorted_categories = []
        for name in names:
            category = self.categories.get(name)
            size = 0
            if category:
                size = category.size
            total_size += size
            if size > 0:
                sorted_categories.append(
                    (name, size, (float(size) * 100) / section_size))
        sorted_categories.sort(key=lambda entry: entry[1], reverse=True)
        for category in sorted_categories:
            if useCSV:
                print("%s;%d;%.2f%%" %
                      (category[0], category[1], category[2]))
            else:
                print("%60s: %8d (%6.2f%%)" %
                      (category[0], category[1], category[2]))
        print("%60s: %8d (%6.2f%%)" % ('TOTAL', total_size, float(100)))

    def uncatorizedSymbols(self):
        category = self.categories.get('Unknown')
        if category:
            return category.symbols
        return None

    def print_uncategorizedSymbols(self):
        syms = self.uncatorizedSymbols()
        if syms:
            for symbol in syms:
                print(symbol.mangled_name + " " + symbol.name + " " +
                      str(symbol.size))

    def print_category(self, category):
        category = self.categories.get(category)
        if category:
            if category.symbols:
                sorted_symbols = sorted(category.symbols, key=get_symbol_size)
                for sym in sorted_symbols:
                    print('%8d %s %s' % (sym.size, sym.name, sym.mangled_name))

    def has_category(self, category):
        category = self.categories.get(category)
        if category:
            if category.symbols:
                return True
        return False

def parse_segments(path, arch):
    mangled = subprocess.check_output(
        ['symbols', '-noSources', '-noDemangling', '-arch', arch, path])
    demangle = subprocess.Popen(
        ['xcrun', 'swift-demangle'], stdin=subprocess.PIPE,
        stdout=subprocess.PIPE)
    demangled = demangle.communicate(mangled)[0]
    symbols = {}
    segments = []
    segment_regex = re.compile(
        r"^        0x[0-9a-f]+ \(\s*0x(?P<size>[0-9a-f]+)\) "
        r"(?P<name>.+?) (?P<name2>.+?)$")
    object_file_segment_regex = re.compile(
        r"^        0x[0-9a-f]+ \(\s*0x(?P<size>[0-9a-f]+)\)  "
        r"SEGMENT$")
    section_regex = re.compile(
        r"^            0x[0-9a-f]+ \(\s*0x(?P<size>[0-9a-f]+)\) "
        r"(?P<name>.+?) (?P<name2>.+?)$")
    symbol_regex = re.compile(
        r"^                0x[0-9a-f]+ \(\s*0x(?P<size>[0-9a-f]+)\) "
        r"(?P<name>.+?) \[[^\]]+\] $")

    mangled_lines = mangled.splitlines()
    current_line_number = 0

    for line in demangled.splitlines():
        mangled_line = mangled_lines[current_line_number]
        current_line_number += 1

        # Match a segment entry.
        segment_match = segment_regex.match(line)
        if segment_match:
            new_segment = Segment(segment_match.group('name'))
            segments.append(new_segment)
            continue

        object_file_segment_match = object_file_segment_regex.match(line)
        if object_file_segment_match:
            new_segment = Segment("SEGMENT")
            segments.append(new_segment)
            continue

        # Match a section entry.
        section_match = section_regex.match(line)
        if section_match:
            new_section = Section(section_match.group('name2'),
                                  int(section_match.group('size'), 16))
            segments[-1].sections.append(new_section)
            continue

        # Match a symbol entry.
        symbol_match = symbol_regex.match(line)
        if not symbol_match:
            continue
        mangled_symbol_match = symbol_regex.match(mangled_line)
        if not mangled_symbol_match:
            print('mangled and demangled mismatch')
            print(mangled_line)
            print(line)
            assert False

        symbol = Symbol(symbol_match.group('name'),
                        mangled_symbol_match.group('name'),
                        int(symbol_match.group('size'), 16))
        existing = symbols.get(symbol.name)
        if existing:
            existing.size += symbol.size
        else:
            symbols[symbol.name] = symbol
            segments[-1].sections[-1].symbols.append(symbol)

    return segments

def show_all(segments):
    for segment in segments:
        for section in segment.sections:
            symbols = section.symbols
            for sym in symbols:
                print(str(sym.size) + ' ' + sym.name + ' ' + sym.mangled_name)

def categorize(segments):
    for segment in segments:
        for section in segment.sections:
            print('Section %52s: %8d' %
                  (segment.name + ';' + section.name, section.size))
            symbols = section.symbols
            categories = Categories()
            categories.categorize(symbols)
            categories.print_summary(section.size)
            print('')

def uncategorized(segments):
    for segment in segments:
        for section in segment.sections:
            symbols = section.symbols
            categories = Categories()
            categories.categorize(symbols)
            categories.print_uncategorizedSymbols()

def list_category(segments, category):
    for segment in segments:
        for section in segment.sections:
            symbols = section.symbols
            categories = Categories()
            categories.categorize(symbols)
            if categories.has_category(category):
                print('Section %22s: %8d' %
                      (segment.name + ';' + section.name, section.size))
                categories.print_category(category)

if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))
<commit_msg>utils: change handling of generic specializations in analyze_code_size.py<commit_after>#!/usr/bin/env python

import argparse
import re
import subprocess
import sys

useCSV = False

def main(arguments):
    parser = argparse.ArgumentParser(
        description='Analyze the code size in a binary')
    parser.add_argument('-arch', type=str,
                        help='the arch to look at', default='arm64')
    parser.add_argument('-categorize', action='store_true',
                        help='categorize symbols', dest='build_categories',
                        default=False)
    parser.add_argument('-list-category', type=str,
                        help='list symbols in category')
    parser.add_argument('-csv', dest='use_csv', action='store_true',
                        help='print results as csv')
    parser.add_argument('-uncategorized', action='store_true',
                        help='show all uncategorized symbols',
                        dest='show_uncategorized',
                        default=False)
    parser.add_argument('bin', help='the binary')
    parser.set_defaults(use_csv=False)

    args = parser.parse_args(arguments)
    if args.use_csv:
        global useCSV
        useCSV = True
        print("Using csv")

    segments = parse_segments(args.bin, args.arch)

    if args.build_categories:
        categorize(segments)
    elif args.show_uncategorized:
        uncategorized(segments)
    elif args.list_category:
        list_category(segments, args.list_category)
    else:
        show_all(segments)

class Symbol(object):
    def __init__(self, name, mangled_name, size):
        self.name = name
        self.mangled_name = mangled_name
        self.count = 1
        self.size = int(size)

def get_symbol_size(sym):
    return sym.size

class Segment(object):
    def __init__(self, name):
        self.name = name
        self.sections = []

class Section(object):
    def __init__(self, name, size):
        self.name = name
        self.size = size
        self.symbols = []

class Category(object):
    def __init__(self, name):
        self.name = name
        self.size = 0
        self.symbols = []

    def add(self, symbol):
        self.symbols.append(symbol)
        self.size += symbol.size

class Categories(object):
    def __init__(self):
        self.category_matching = [
            ['Objective-C function', re.compile(r'.*[+-]\[')],
            ['C++', re.compile(r'_+swift')],
            ['Generic specialization of stdlib', re.compile(r'.*generic specialization.* of Swift\.')],
            ['Generic specialization', re.compile(r'.*generic specialization')],
            ['Merged function', re.compile(r'merged ')],
            ['Key path', re.compile(r'key path')],
            ['Function signature specialization',
                re.compile(r'function signature specialization')],
            ['Reabstraction thunk helper',
                re.compile(r'reabstraction thunk helper')],
            ['vtable thunk', re.compile(r'vtable thunk for')],
            ['@objc thunk', re.compile(r'@objc')],
            ['@nonobjc thunk', re.compile(r'@nonobjc')],
            ['Value witness', re.compile(r'.*value witness for')],
            ['Block copy helper', re.compile(r'_block_copy_helper')],
            ['Block destroy helper', re.compile(r'_block_destroy_helper')],
            ['Block literal global', re.compile(r'___block_literal_global')],
            ['Destroy helper block', re.compile(r'___destroy_helper_block')],
            ['Copy helper block', re.compile(r'___copy_helper_block')],
            ['Object destroy', re.compile(r'_objectdestroy')],
            ['Partial apply forwarder',
                re.compile(r'partial apply forwarder')],
            ['Closure function', re.compile(r'closure #')],
            ['ObjC metadata update function',
                re.compile(r'ObjC metadata update function for')],
            ['Variable initialization expression',
                re.compile(r'variable initialization expression of')],
            ['Global initialization', re.compile(r'_globalinit_')],
            ['Unnamed', re.compile(r'___unnamed_')],
            ['Dyld stubs', re.compile(r'DYLD-STUB\$')],
            ['Witness table accessor',
                re.compile(r'.*witness table accessor for')],
            ['Protocol witness', re.compile(r'protocol witness for')],
            ['Outlined variable', re.compile(r'outlined variable #')],
            ['Outlined value function (copy,destroy,release...)',
                re.compile(r'outlined')],
            ['_symbolic', re.compile(r'_symbolic')],
            ['_associated conformance',
                re.compile(r'_associated conformance')],
            ['Direct field offset', re.compile(r'direct field offset for')],
            ['Value witness tables', re.compile(r'.*value witness table')],
            ['Protocol witness table',
                re.compile(r'.*protocol witness table for')],
            ['Protocol conformance descriptor',
                re.compile(r'protocol conformance descriptor for')],
            ['Lazy protocol witness table cache var',
                re.compile(
                    r'lazy protocol witness table cache variable for type')],
            ['Nominal type descriptor',
                re.compile(r'nominal type descriptor for')],
            ['ObjC class', re.compile(r'_OBJC_CLASS_')],
            ['ObjC metaclass', re.compile(r'_OBJC_METACLASS')],
            ['ObjC ivar', re.compile(r'_OBJC_IVAR')],
            ['Metaclass', re.compile(r'metaclass for')],
            ['Block descriptor', re.compile(r'_+block_descriptor')],
            ['Extension descriptor', re.compile(r'extension descriptor')],
            ['Module descriptor', re.compile(r'module descriptor')],
            ['Associated type descriptor',
                re.compile(r'associated type descriptor for')],
            ['Associated conformance descriptor',
                re.compile(r'associated conformance descriptor for')],
            ['Protocol descriptor', re.compile(r'protocol descriptor for')],
            ['Base conformance descriptor',
                re.compile(r'base conformance descriptor for')],
            ['Protocol requirements base descriptor',
                re.compile(r'protocol requirements base descriptor for')],
            ['Property descriptor', re.compile(r'property descriptor for')],
            ['Method descriptor', re.compile(r'method descriptor for')],
            ['Anonymous descriptor', re.compile(r'anonymous descriptor')],
            ['Type metadata accessor',
                re.compile(r'.*type metadata accessor')],
            ['Type metadata', re.compile(r'.*type metadata')],
            ['Reflection metadata descriptor',
                re.compile(r'reflection metadata .* descriptor')],
        ]

        self.category_mangled_matching = [
            ['Swift variable storage', re.compile(r'^_\$s.*[v][p][Z]?$')],
            ['Swift constructor', re.compile(r'^_\$s.*[f][cC]$')],
            ['Swift initializer', re.compile(r'^_\$s.*[f][ie]$')],
            ['Swift destructor/destroyer', re.compile(r'^_\$s.*[f][dDE]$')],
            ['Swift getter', re.compile(r'^_\$s.*[iv][gG]$')],
            ['Swift setter', re.compile(r'^_\$s.*[iv][swW]$')],
            ['Swift materializeForSet', re.compile(r'^_\$s.*[iv][m]$')],
            ['Swift modify', re.compile(r'^_\$s.*[iv][M]$')],
            ['Swift read', re.compile(r'^_\$s.*[iv][r]$')],
            ['Swift addressor', re.compile(r'^_\$s.*[iv][al][uOop]$')],
            ['Swift function', re.compile(r'^_\$s.*F$')],
            ['Swift unknown', re.compile(r'^_\$s.*')],
        ]
        self.categories = {}

    def categorize_by_name(self, symbol):
        for c in self.category_matching:
            if c[1].match(symbol.name):
                return c[0]
        return None

    def categorize_by_mangled_name(self, symbol):
        for c in self.category_mangled_matching:
            if c[1].match(symbol.mangled_name):
                return c[0]
        return None

    def add_symbol(self, category_name, symbol):
        existing_category = self.categories.get(category_name)
        if existing_category:
            existing_category.add(symbol)
        else:
            new_category = Category(category_name)
            new_category.add(symbol)
            self.categories[category_name] = new_category

    def add(self, symbol):
        category_name = self.categorize_by_name(symbol)
        if category_name:
            self.add_symbol(category_name, symbol)
            return
        category_name = self.categorize_by_mangled_name(symbol)
        if category_name:
            self.add_symbol(category_name, symbol)
        else:
            self.add_symbol('Unknown', symbol)

    def categorize(self, symbols):
        for sym in symbols:
            self.add(sym)

    def print_summary(self, section_size):
        names = [c[0] for c in self.category_matching]
        names.extend([c[0] for c in self.category_mangled_matching])
        names.append('Unknown')
        total_size = 0
        sorted_categories = []
        for name in names:
            category = self.categories.get(name)
            size = 0
            if category:
                size = category.size
            total_size += size
            if size > 0:
                sorted_categories.append(
                    (name, size, (float(size) * 100) / section_size))
        sorted_categories.sort(key=lambda entry: entry[1], reverse=True)
        for category in sorted_categories:
            if useCSV:
                print("%s;%d;%.2f%%" %
                      (category[0], category[1], category[2]))
            else:
                print("%60s: %8d (%6.2f%%)" %
                      (category[0], category[1], category[2]))
        print("%60s: %8d (%6.2f%%)" % ('TOTAL', total_size, float(100)))

    def uncatorizedSymbols(self):
        category = self.categories.get('Unknown')
        if category:
            return category.symbols
        return None

    def print_uncategorizedSymbols(self):
        syms = self.uncatorizedSymbols()
        if syms:
            for symbol in syms:
                print(symbol.mangled_name + " " + symbol.name + " " +
                      str(symbol.size))

    def print_category(self, category):
        category = self.categories.get(category)
        if category:
            if category.symbols:
                sorted_symbols = sorted(category.symbols, key=get_symbol_size)
                for sym in sorted_symbols:
                    print('%8d %s %s' % (sym.size, sym.name, sym.mangled_name))

    def has_category(self, category):
        category = self.categories.get(category)
        if category:
            if category.symbols:
                return True
        return False

def parse_segments(path, arch):
    mangled = subprocess.check_output(
        ['symbols', '-noSources', '-noDemangling', '-arch', arch, path])
    demangle = subprocess.Popen(
        ['xcrun', 'swift-demangle'], stdin=subprocess.PIPE,
        stdout=subprocess.PIPE)
    demangled = demangle.communicate(mangled)[0]
    symbols = {}
    segments = []
    segment_regex = re.compile(
        r"^        0x[0-9a-f]+ \(\s*0x(?P<size>[0-9a-f]+)\) "
        r"(?P<name>.+?) (?P<name2>.+?)$")
    object_file_segment_regex = re.compile(
        r"^        0x[0-9a-f]+ \(\s*0x(?P<size>[0-9a-f]+)\)  "
        r"SEGMENT$")
    section_regex = re.compile(
        r"^            0x[0-9a-f]+ \(\s*0x(?P<size>[0-9a-f]+)\) "
        r"(?P<name>.+?) (?P<name2>.+?)$")
    symbol_regex = re.compile(
        r"^                0x[0-9a-f]+ \(\s*0x(?P<size>[0-9a-f]+)\) "
        r"(?P<name>.+?) \[[^\]]+\] $")

    mangled_lines = mangled.splitlines()
    current_line_number = 0

    for line in demangled.splitlines():
        mangled_line = mangled_lines[current_line_number]
        current_line_number += 1

        # Match a segment entry.
        segment_match = segment_regex.match(line)
        if segment_match:
            new_segment = Segment(segment_match.group('name'))
            segments.append(new_segment)
            continue

        object_file_segment_match = object_file_segment_regex.match(line)
        if object_file_segment_match:
            new_segment = Segment("SEGMENT")
            segments.append(new_segment)
            continue

        # Match a section entry.
        section_match = section_regex.match(line)
        if section_match:
            new_section = Section(section_match.group('name2'),
                                  int(section_match.group('size'), 16))
            segments[-1].sections.append(new_section)
            continue

        # Match a symbol entry.
        symbol_match = symbol_regex.match(line)
        if not symbol_match:
            continue
        mangled_symbol_match = symbol_regex.match(mangled_line)
        if not mangled_symbol_match:
            print('mangled and demangled mismatch')
            print(mangled_line)
            print(line)
            assert False

        symbol = Symbol(symbol_match.group('name'),
                        mangled_symbol_match.group('name'),
                        int(symbol_match.group('size'), 16))
        existing = symbols.get(symbol.name)
        if existing:
            existing.size += symbol.size
        else:
            symbols[symbol.name] = symbol
            segments[-1].sections[-1].symbols.append(symbol)

    return segments

def show_all(segments):
    for segment in segments:
        for section in segment.sections:
            symbols = section.symbols
            for sym in symbols:
                print(str(sym.size) + ' ' + sym.name + ' ' + sym.mangled_name)

def categorize(segments):
    for segment in segments:
        for section in segment.sections:
            print('Section %52s: %8d' %
                  (segment.name + ';' + section.name, section.size))
            symbols = section.symbols
            categories = Categories()
            categories.categorize(symbols)
            categories.print_summary(section.size)
            print('')

def uncategorized(segments):
    for segment in segments:
        for section in segment.sections:
            symbols = section.symbols
            categories = Categories()
            categories.categorize(symbols)
            categories.print_uncategorizedSymbols()

def list_category(segments, category):
    for segment in segments:
        for section in segment.sections:
            symbols = section.symbols
            categories = Categories()
            categories.categorize(symbols)
            if categories.has_category(category):
                print('Section %22s: %8d' %
                      (segment.name + ';' + section.name, section.size))
                categories.print_category(category)

if __name__ == '__main__':
    sys.exit(main(sys.argv[1:]))
<|endoftext|>


<commit_before>#!/usr/bin/python
import argparse
import os
import sys
import subprocess

# write func to find files

class concatenateTrajectories:
	def __init__(self, folder, trajectoryFileType, checkConcatGroup, systemName, repetition):
		# Check if the correct file type was chosen
		self.trajectoryFileType = self.checkFileType(trajectoryFileType)

		# Check if selected group is correct
		self.concatGroup = self.checkConcatGroup(checkConcatGroup)

		# Get all the trajectory names
		self.trajectories = self.checkCorrectFolder(folder, trajectoryFileType)

		# Get the individual trajectory length
		self.trajectoriesLength = self.findTrajectoryLength()

		# System name
		self.systemName = systemName

		# Repetition
		self.repetition = repetition

		# TPR file for this series
		self.tprFile = ""

	def findFilesOfType(self, path, filetype):
		foundFiles = []

		# If path is necessary check if there is a final /
		if not path == "":
			if path[-1] != "/":
				path = "%s/"%path

		files = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]

		# Loop over all path
		for file in files:
			# IF path is necessary add it to file name
			if not path == "":
				fileAndPath = path+file
			# only show file of the filetype chosen and which do not start with a .
			if fileAndPath.endswith(".%s"%filetype) and file[0] != '.':
				foundFiles.append(fileAndPath)

		# Return files found
		return foundFiles

	def concatenate(self):
		# To concatenate the trajectories we need to find the timestamps of the trajectories
		# It starts at 0
		timeStamp = [0]
		for currentLength in self.trajectoriesLength:
			timeStamp.append(timeStamp[-1]+currentLength)

		# We do not need the last time stamp
		del timeStamp[-1]

		# Time stamp string for the GROMACS concatenate tool
		timeStampString = '\n'.join(str(e) for e in timeStamp)

		# Trjactory path string for the GROMACS concatenate tool
		trajectoryPath = ' '.join(self.trajectories)

		# Find the tpr file
		tprFiles = self.findFilesOfType(os.path.split(self.trajectories[0])[0], "tpr")
		if len(tprFiles) != 1:
			sys.exit("Found %s tpr files in folder %s. Please make sure only the correct tpr file is present")
		else:
			self.tprFile = tprFiles[0]

		# Generate an index file for the concatenate tool to ask for the group to use in the new trajectory
		makeNDXProgram = self.which("make_ndx_mpi")

		# Generate a logfile
		logFile = open("%s_concatenate_rep%i.log"%(self.systemName, self.repetition), 'a')

		# Run the program
		p = subprocess.Popen("echo \"q\" | %s -f %s -o %s_rep%i_protein.ndx"%(makeNDXProgram, self.tprFile, self.systemName, self.repetition), shell=True, stdin=subprocess.PIPE, stdout=logFile, stderr=subprocess.STDOUT)
		p.communicate() #now wait

		# Generate the concatenate trajectory
		gmxProgram = self.which("gmx_mpi")

		# Time stamp string for the GROMACS concatenate tool
		timeStampString = "0\n"
		for length in self.trajectoriesLength:
			timeStampString += "%i\n"%(int(timeStampString.split()[-1])+length)

		# Trjactory path string for the GROMACS concatenate tool
		trajectoryPath = ' '.join(self.trajectories)

		p = subprocess.Popen("echo \"1 %s\" | %s trjcat -f %s -settime -o %s_0-%i_rep%i.xtc -n %s_rep%i_protein.ndx"%(timeStampString, gmxProgram, trajectoryPath, self.systemName, sum(self.trajectoriesLength)/1000, self.repetition, self.systemName, self.repetition), shell=True, stdin=subprocess.PIPE, stdout=logFile, stderr=subprocess.STDOUT)
		p.communicate() #now wait

	def which(self, program):
		def is_exe(fpath):
			return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

		fpath, fname = os.path.split(program)
		if fpath:
			if is_exe(program):
				return program
			elif is_exe("%s_mpi"%program):
				return "%s_mpi"%program
			elif is_exe(program[:len(program)-4]):
				return program[:len(program)-4]
		else:
			for path in os.environ["PATH"].split(os.pathsep):
				path = path.strip('"')
				exe_file = os.path.join(path, program)
				if is_exe(exe_file):
					return exe_file
				elif is_exe("%s_mpi"%exe_file):
					return "%s_mpi"%exe_file
				elif is_exe(exe_file[:len(exe_file)-4]):
					return exe_file[:len(exe_file)-4]
		return False

	def checkConcatGroup(self, group):
		# Define the supported groups
		supportedGroupTypes = {"system" : 0, "protein" : 1, "non-water" : 14}

		if group.lower() in supportedGroupTypes:
			return supportedGroupTypes[group.lower()]
		else:
			sys.exit("The group %s is not supported"%group)

	def checkFileType(self, filetype):
		# Define the supported file types
		supportedFileTypes = ["xtc", "trr"]

		# Check if the file type is supported
		if filetype in supportedFileTypes:
			return filetype
		else:
			sys.exit("The file type %s is not supported"%filetype)

	def checkCorrectFolder(self, folder, filetype):
		# Check if the chosen folder exists
		if not os.path.exists(folder):
			sys.exit("%s does not exist"%folder)

		trj_files = []

		# Look for trajectory files
		for root, dirs, files in os.walk(folder):
			foundTrjFile = self.findFilesOfType(root, filetype)

			# If we find more than one file stop
			if len(foundTrjFile) > 1:
				sys.exit("There are multiple %s files in the folder %s"%(filetype, root))

			# If there is no file in the folder continue
			elif len(foundTrjFile) == 0:
				continue
			# If there is one file add it to the trajectories
			trj_files.append(foundTrjFile[0])
		# Verify if we found some trajectories
		if len(trj_files) == 0:
			sys.exit("No trajectory files of type %s in %s and it's subfolders found"%(filetype, folder))
		return trj_files

	def findTrajectoryLength(self):
		# Find the correcy log file and the length in the log file
		trj_length = []

		# Find all log files in all trajectory folders
		for trajectory in self.trajectories:
			trajectoryPath = os.path.split(trajectory)[0]

			logFiles = self.findFilesOfType(trajectoryPath, "log")

			# Search for the correct log file
			trj_logFile = ""
			for logFile in logFiles:
				with open(logFile, 'r') as inF:
					for line in inF:
						# Search for the first line in a standard trajectory output
						if "Log file opened on" in line:
							trj_logFile = logFile
							break
			# if the standard trajectory output was not found no log file for the trajectory
			if trj_logFile == "":
				sys.exit("Could not find the log file for %s"%trajectory)

			# If we found the trajectory extract the length dt*nsteps
			else:
				with open(logFile, 'r') as inF:
					for line in inF:
						if ' dt ' in line:
							sizeOfSteps = float(line.split()[2])
						if 'nsteps' in line:
							numberOfSteps = int(line.split()[2])
							break
					trj_length.append(int(sizeOfSteps*numberOfSteps))
		return trj_length

##### GETTER ######

	# Get the trajectory length and name
	def getTrajLength(self):
		trajLength = []
		for index, value in enumerate(self.trajectories):
			trajLength.append([value, self.trajectoriesLength[index]])
		return trajLength

if __name__ == "__main__":

	#####################################
	#									#
	# The help text						#
	#									#
	#####################################

	parser = argparse.ArgumentParser(description="""
	g_trajAnal
	Maximilian Ebert 2016

	This tool let's you analyse and concatenate GROMACS trajectories based
	on the data structure described in my GROMACS simulation tutorial. You
	must follow the same folder and file nomenclature so that this script
	works out of the box. Otherwise you have to tinker with the code.

	""", formatter_class=argparse.RawDescriptionHelpFormatter)

	# Definition of arguments

	# Trajectory file or folder holding sub directories with trajectories
	parser.add_argument("-cf", metavar="[<folder to .xtc/.trr>] (Input)", help="Concatenate trajectories in folder", required=True)
	parser.add_argument("-ct", metavar="(xtc) ", help="Concatenate input file type: xtc, trr", default="xtc")
	parser.add_argument("-cg", metavar="(protein) ", help="Concatenate groupe: System, Protein, non-Water", default="Protein")
	parser.add_argument("-n", metavar="", help="Name of the system (ex. 1XPB)", required=True)
	parser.add_argument("-r", metavar="(1) ", help="Repetition of run", default=1)
	args = parser.parse_args()

	trj = concatenateTrajectories(args.cf, args.ct, args.cg, args.n, args.r)
	trj.concatenate()

<commit_msg>Change group def from init to concat function<commit_after>#!/usr/bin/python
import argparse
import os
import sys
import subprocess

# write func to find files

class concatenateTrajectories:
	def __init__(self, folder, trajectoryFileType, systemName, repetition):
		# Check if the correct file type was chosen
		self.trajectoryFileType = self.checkFileType(trajectoryFileType)

		# Get all the trajectory names
		self.trajectories = self.checkCorrectFolder(folder, trajectoryFileType)

		# Get the individual trajectory length
		self.trajectoriesLength = self.findTrajectoryLength()

		# System name
		self.systemName = systemName

		# Repetition
		self.repetition = repetition

		# TPR file for this series
		self.tprFile = ""

	def findFilesOfType(self, path, filetype):
		foundFiles = []

		# If path is necessary check if there is a final /
		if not path == "":
			if path[-1] != "/":
				path = "%s/"%path

		files = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]

		# Loop over all path
		for file in files:
			# IF path is necessary add it to file name
			if not path == "":
				fileAndPath = path+file
			# only show file of the filetype chosen and which do not start with a .
			if fileAndPath.endswith(".%s"%filetype) and file[0] != '.':
				foundFiles.append(fileAndPath)

		# Return files found
		return foundFiles

	def concatenate(self, checkConcatGroup):
		# Check if selected group is correct
		concatGroup = self.checkConcatGroup(checkConcatGroup)

		# To concatenate the trajectories we need to find the timestamps of the trajectories
		# It starts at 0
		timeStamp = [0]
		for currentLength in self.trajectoriesLength:
			timeStamp.append(timeStamp[-1]+currentLength)

		# We do not need the last time stamp
		del timeStamp[-1]

		# Time stamp string for the GROMACS concatenate tool
		timeStampString = '\n'.join(str(e) for e in timeStamp)

		# Trjactory path string for the GROMACS concatenate tool
		trajectoryPath = ' '.join(self.trajectories)

		# Find the tpr file
		tprFiles = self.findFilesOfType(os.path.split(self.trajectories[0])[0], "tpr")
		if len(tprFiles) != 1:
			sys.exit("Found %s tpr files in folder %s. Please make sure only the correct tpr file is present")
		else:
			self.tprFile = tprFiles[0]

		# Generate an index file for the concatenate tool to ask for the group to use in the new trajectory
		makeNDXProgram = self.which("make_ndx_mpi")

		# Generate a logfile
		logFile = open("%s_concatenate_rep%i.log"%(self.systemName, self.repetition), 'a')

		# Run the program
		p = subprocess.Popen("echo \"q\" | %s -f %s -o %s_rep%i_protein.ndx"%(makeNDXProgram, self.tprFile, self.systemName, self.repetition), shell=True, stdin=subprocess.PIPE, stdout=logFile, stderr=subprocess.STDOUT)
		p.communicate() #now wait

		# Generate the concatenate trajectory
		gmxProgram = self.which("gmx_mpi")

		# Time stamp string for the GROMACS concatenate tool
		timeStampString = "0\n"
		for length in self.trajectoriesLength:
			timeStampString += "%i\n"%(int(timeStampString.split()[-1])+length)

		# Trjactory path string for the GROMACS concatenate tool
		trajectoryPath = ' '.join(self.trajectories)

		p = subprocess.Popen("echo \"1 %s\" | %s trjcat -f %s -settime -o %s_0-%i_rep%i.xtc -n %s_rep%i_protein.ndx"%(timeStampString, gmxProgram, trajectoryPath, self.systemName, sum(self.trajectoriesLength)/1000, self.repetition, self.systemName, self.repetition), shell=True, stdin=subprocess.PIPE, stdout=logFile, stderr=subprocess.STDOUT)
		p.communicate() #now wait

	def which(self, program):
		def is_exe(fpath):
			return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

		fpath, fname = os.path.split(program)
		if fpath:
			if is_exe(program):
				return program
			elif is_exe("%s_mpi"%program):
				return "%s_mpi"%program
			elif is_exe(program[:len(program)-4]):
				return program[:len(program)-4]
		else:
			for path in os.environ["PATH"].split(os.pathsep):
				path = path.strip('"')
				exe_file = os.path.join(path, program)
				if is_exe(exe_file):
					return exe_file
				elif is_exe("%s_mpi"%exe_file):
					return "%s_mpi"%exe_file
				elif is_exe(exe_file[:len(exe_file)-4]):
					return exe_file[:len(exe_file)-4]
		return False

	def checkConcatGroup(self, group):
		# Define the supported groups
		supportedGroupTypes = {"system" : 0, "protein" : 1, "non-water" : 14}

		if group.lower() in supportedGroupTypes:
			return supportedGroupTypes[group.lower()]
		else:
			sys.exit("The group %s is not supported"%group)

	def checkFileType(self, filetype):
		# Define the supported file types
		supportedFileTypes = ["xtc", "trr"]

		# Check if the file type is supported
		if filetype in supportedFileTypes:
			return filetype
		else:
			sys.exit("The file type %s is not supported"%filetype)

	def checkCorrectFolder(self, folder, filetype):
		# Check if the chosen folder exists
		if not os.path.exists(folder):
			sys.exit("%s does not exist"%folder)

		trj_files = []

		# Look for trajectory files
		for root, dirs, files in os.walk(folder):
			foundTrjFile = self.findFilesOfType(root, filetype)

			# If we find more than one file stop
			if len(foundTrjFile) > 1:
				sys.exit("There are multiple %s files in the folder %s"%(filetype, root))

			# If there is no file in the folder continue
			elif len(foundTrjFile) == 0:
				continue
			# If there is one file add it to the trajectories
			trj_files.append(foundTrjFile[0])
		# Verify if we found some trajectories
		if len(trj_files) == 0:
			sys.exit("No trajectory files of type %s in %s and it's subfolders found"%(filetype, folder))
		return trj_files

	def findTrajectoryLength(self):
		# Find the correcy log file and the length in the log file
		trj_length = []

		# Find all log files in all trajectory folders
		for trajectory in self.trajectories:
			trajectoryPath = os.path.split(trajectory)[0]

			logFiles = self.findFilesOfType(trajectoryPath, "log")

			# Search for the correct log file
			trj_logFile = ""
			for logFile in logFiles:
				with open(logFile, 'r') as inF:
					for line in inF:
						# Search for the first line in a standard trajectory output
						if "Log file opened on" in line:
							trj_logFile = logFile
							break
			# if the standard trajectory output was not found no log file for the trajectory
			if trj_logFile == "":
				sys.exit("Could not find the log file for %s"%trajectory)

			# If we found the trajectory extract the length dt*nsteps
			else:
				with open(logFile, 'r') as inF:
					for line in inF:
						if ' dt ' in line:
							sizeOfSteps = float(line.split()[2])
						if 'nsteps' in line:
							numberOfSteps = int(line.split()[2])
							break
					trj_length.append(int(sizeOfSteps*numberOfSteps))
		return trj_length

##### GETTER ######

	# Get the trajectory length and name
	def getTrajLength(self):
		trajLength = []
		for index, value in enumerate(self.trajectories):
			trajLength.append([value, self.trajectoriesLength[index]])
		return trajLength

if __name__ == "__main__":

	#####################################
	#									#
	# The help text						#
	#									#
	#####################################

	parser = argparse.ArgumentParser(description="""
	g_trajAnal
	Maximilian Ebert 2016

	This tool let's you analyse and concatenate GROMACS trajectories based
	on the data structure described in my GROMACS simulation tutorial. You
	must follow the same folder and file nomenclature so that this script
	works out of the box. Otherwise you have to tinker with the code.

	""", formatter_class=argparse.RawDescriptionHelpFormatter)

	# Definition of arguments

	# Trajectory file or folder holding sub directories with trajectories
	parser.add_argument("-cf", metavar="[<folder to .xtc/.trr>] (Input)", help="Concatenate trajectories in folder", required=True)
	parser.add_argument("-ct", metavar="(xtc) ", help="Concatenate input file type: xtc, trr", default="xtc")
	parser.add_argument("-cg", metavar="(protein) ", help="Concatenate groupe: System, Protein, non-Water", default="Protein")
	parser.add_argument("-n", metavar="", help="Name of the system (ex. 1XPB)", required=True)
	parser.add_argument("-r", metavar="(1) ", help="Repetition of run", default=1)
	args = parser.parse_args()

	# Initiate the concatenate trajectory object
	trj = concatenateTrajectories(args.cf, args.ct, args.n, args.r)

	# Concatenate
	trj.concatenate(args.cg)

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging
import itertools

import requests
from lxml import html
from dateutil.parser import parse

from db.common import session_scope
from db.team import Team
from db.player import Player
from db.player_data_item import PlayerDataItem

logger = logging.getLogger(__name__)

CAPFRIENDLY_PLAYER_PREFIX = "http://www.capfriendly.com/players/"
CAPFRIENDLY_TEAM_PREFIX = "http://www.capfriendly.com/teams/"

def retrieve_capfriendly_ids(team_id):
    """
    Retrieves ids from capfriendly.com for all players of the team with
    the specified id
    """
    team = Team.find_by_id(team_id)

    logger.info(
        "+ Retrieving capfriendly ids for all players of the %s" % team)

    url = "".join((
        CAPFRIENDLY_TEAM_PREFIX,
        team.team_name.replace(" ", "").lower()))

    r = requests.get(url)
    doc = html.fromstring(r.text)

    player_name_trs = doc.xpath("//tr[@class='even c' or @class='odd c']")

    for tr in player_name_trs:
        player_name = tr.xpath("td/a/text()").pop(0)
        player_link = tr.xpath("td/a/@href").pop(0)
        player_position = tr.xpath("td[3]/span/text()")
        if player_position:
            player_position = player_position.pop(0)
        else:
            player_position = ''

        last_name, first_name = player_name.split(", ")
        plr = Player.find_by_name_extended(first_name, last_name)
        if plr is None and player_position:
            primary_position = player_position.split(", ")[0][0]
            plr = Player.find_by_name_position(
                first_name, last_name, primary_position)
        if plr and plr.capfriendly_id is None:
            print(
                "+ Ambigious capfriendly id for player:", first_name,
                last_name, plr, player_link.split("/")[-1])
        if plr is None:
            print(
                "+ No player for capfriendly id:", first_name,
                last_name, player_link.split("/")[-1])

def retrieve_capfriendly_id(player_id):
    """
    Retrieves an id from capfriendly.com for the player with the
    specified id.
    """
    plr = Player.find_by_id(player_id)
    pdi = PlayerDataItem.find_by_player_id(player_id)

    if plr.capfriendly_id is not None:
        logger.info(
            "+ Existing capfriendly id for %s: %s" % (
                plr.name, plr.capfriendly_id))
        return plr.capfriendly_id

    # compiling all potential capfriendly ids from the player's name(s)
    potential_capfriendly_ids = collect_potential_capfriendly_ids(plr)

    capfriendly_id_found = False

    while potential_capfriendly_ids and not capfriendly_id_found:
        potential_capfriendly_id = potential_capfriendly_ids.pop(0)
        query_id = potential_capfriendly_id.replace(" ", "-")
        url = "".join((CAPFRIENDLY_PLAYER_PREFIX, query_id))
        req = requests.get(url)
        doc = html.fromstring(req.text)
        # retrieving player name from capfriendly page
        page_header = doc.xpath("//h1/text()").pop(0).strip()
        # retrieving player's date of birth from capfriendly page
        page_dob = doc.xpath(
            "//span[@class='l pld_l']/ancestor::div/text()")[0].strip()
        page_dob = parse(page_dob).date()
        # comparing names
        if page_header == potential_capfriendly_id.upper():
            # comparing date of births
            if page_dob == pdi.date_of_birth:
                capfriendly_id_found = True
                logger.info(
                    "+ Found capfriendly id for %s: %s" % (plr.name, query_id))
                plr.capfriendly_id = query_id
                with session_scope() as session:
                    session.merge(plr)
                    session.commit()

    if not capfriendly_id_found:
        logger.warn("+ No capfriendly id found for %s" % plr.name)

    return plr.capfriendly_id

def collect_potential_capfriendly_ids(plr):
    """
    Compiles all potential combinations of player first and last names
    to find a potential capfriendly id.
    """
    # listing all of players' potential first names
    first_names = [plr.first_name]
    if plr.alternate_first_names:
        first_names += plr.alternate_first_names
    first_names = list(map(str.lower, first_names))

    # listing all of players' potential last names
    last_names = [plr.last_name]
    if plr.alternate_last_names:
        last_names += plr.alternate_last_names
    last_names = list(map(str.lower, last_names))

    return list(map(" ".join, itertools.product(first_names, last_names)))
<commit_msg>Disregard dots in potential and retrieved capfriendly ids<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import logging
import itertools

import requests
from lxml import html
from dateutil.parser import parse

from db.common import session_scope
from db.team import Team
from db.player import Player
from db.player_data_item import PlayerDataItem

logger = logging.getLogger(__name__)

CAPFRIENDLY_PLAYER_PREFIX = "http://www.capfriendly.com/players/"
CAPFRIENDLY_TEAM_PREFIX = "http://www.capfriendly.com/teams/"

def retrieve_capfriendly_ids(team_id):
    """
    Retrieves ids from capfriendly.com for all players of the team with
    the specified id
    """
    team = Team.find_by_id(team_id)

    logger.info(
        "+ Retrieving capfriendly ids for all players of the %s" % team)

    url = "".join((
        CAPFRIENDLY_TEAM_PREFIX,
        team.team_name.replace(" ", "").lower()))

    r = requests.get(url)
    doc = html.fromstring(r.text)

    player_name_trs = doc.xpath("//tr[@class='even c' or @class='odd c']")

    for tr in player_name_trs:
        player_name = tr.xpath("td/a/text()").pop(0)
        player_link = tr.xpath("td/a/@href").pop(0)
        player_position = tr.xpath("td[3]/span/text()")
        if player_position:
            player_position = player_position.pop(0)
        else:
            player_position = ''

        last_name, first_name = player_name.split(", ")
        plr = Player.find_by_name_extended(first_name, last_name)
        if plr is None and player_position:
            primary_position = player_position.split(", ")[0][0]
            plr = Player.find_by_name_position(
                first_name, last_name, primary_position)
        if plr and plr.capfriendly_id is None:
            print(
                "+ Ambigious capfriendly id for player:", first_name,
                last_name, plr, player_link.split("/")[-1])
        if plr is None:
            print(
                "+ No player for capfriendly id:", first_name,
                last_name, player_link.split("/")[-1])

def retrieve_capfriendly_id(player_id):
    """
    Retrieves an id from capfriendly.com for the player with the
    specified id.
    """
    plr = Player.find_by_id(player_id)
    pdi = PlayerDataItem.find_by_player_id(player_id)

    if plr.capfriendly_id is not None:
        logger.info(
            "+ Existing capfriendly id for %s: %s" % (
                plr.name, plr.capfriendly_id))
        return plr.capfriendly_id

    # compiling all potential capfriendly ids from the player's name(s)
    potential_capfriendly_ids = collect_potential_capfriendly_ids(plr)

    capfriendly_id_found = False

    while potential_capfriendly_ids and not capfriendly_id_found:
        potential_capfriendly_id = potential_capfriendly_ids.pop(0)
        query_id = potential_capfriendly_id.replace(" ", "-")
        url = "".join((CAPFRIENDLY_PLAYER_PREFIX, query_id))
        req = requests.get(url)
        doc = html.fromstring(req.text)
        # retrieving player name from capfriendly page
        page_header = doc.xpath("//h1/text()").pop(0).strip().replace(".", "")
        # retrieving player's date of birth from capfriendly page
        page_dob = doc.xpath(
            "//span[@class='l pld_l']/ancestor::div/text()")[0].strip()
        page_dob = parse(page_dob).date()
        # comparing names
        if page_header == potential_capfriendly_id.upper().replace(".", ""):
            # comparing date of births
            if page_dob == pdi.date_of_birth:
                capfriendly_id_found = True
                # removing dots from query id to create actual capfriendly id
                found_capfriendly_id = query_id.replace(".", "")
                logger.info(
                    "+ Found capfriendly id for %s: %s" % (
                        plr.name, found_capfriendly_id))
                plr.capfriendly_id = found_capfriendly_id
                with session_scope() as session:
                    session.merge(plr)
                    session.commit()

    if not capfriendly_id_found:
        logger.warn("+ No capfriendly id found for %s" % plr.name)

    return plr.capfriendly_id

def collect_potential_capfriendly_ids(plr):
    """
    Compiles all potential combinations of player first and last names
    to find a potential capfriendly id.
    """
    # listing all of players' potential first names
    first_names = [plr.first_name]
    if plr.alternate_first_names:
        first_names += plr.alternate_first_names
    first_names = list(map(str.lower, first_names))

    # listing all of players' potential last names
    last_names = [plr.last_name]
    if plr.alternate_last_names:
        last_names += plr.alternate_last_names
    last_names = list(map(str.lower, last_names))

    return list(map(" ".join, itertools.product(first_names, last_names)))
<|endoftext|>


<commit_before>from django.conf import settings
from django.forms import widgets
from django.utils.encoding import force_unicode
from django.utils.safestring import mark_safe
from django.forms.util import flatatt
from gmap_admin import settings as gmap_settings

class GoogleMapsWidget(widgets.HiddenInput):

	class Media:
		js = [
			'http://maps.google.com/maps/api/js?sensor=false',
			'gmap_admin/js/google-maps-admin.js',
	        ]

	def render(self, name, value, attrs=None):
		if value is None:
			value = ""
		final_attrs = self.build_attrs(attrs, type=self.input_type, name=name)
		if value != "":
			final_attrs['value'] = force_unicode(self._format_value(value))
		html = u"""
					<input%(attrs)s />

Remove Marker

Double click to zoom in and center on a location. Right click to set the marker on  a position. You can also drag and drop the marker

				""" % {
				        'field_id' : attrs['id'],
				        'delete_id' : 'map_delete_%s'%attrs['id'],
				        'map_id' : 'map_%s'%attrs['id'],
						'attrs' : flatatt(final_attrs),
						'height': gmap_settings.HEIGHT,
						'width': gmap_settings.WIDTH,
						'zoom' : gmap_settings.ZOOM,
						'center_lng' : gmap_settings.DEFAULT_LNG,
						'center_lat' : gmap_settings.DEFAULT_LAT,
					}
		return mark_safe(html)<commit_msg>make map centre on existing point, if there is one, instead of default centre<commit_after>from django.conf import settings
from django.forms import widgets
from django.utils.encoding import force_unicode
from django.utils.safestring import mark_safe
from django.forms.util import flatatt
from gmap_admin import settings as gmap_settings

class GoogleMapsWidget(widgets.HiddenInput):

	class Media:
		js = [
			'http://maps.google.com/maps/api/js?sensor=false',
			'gmap_admin/js/google-maps-admin.js',
			]

	def render(self, name, value, attrs=None):
		if value is None:
			value = ""
			center_lng = gmap_settings.DEFAULT_LNG
			center_lat = gmap_settings.DEFAULT_LAT
		final_attrs = self.build_attrs(attrs, type=self.input_type, name=name)
		if value != "":
			final_attrs['value'] = force_unicode(self._format_value(value))
			center_lat,center_lng = final_attrs['value'].split(',')
		html = u"""
					<input%(attrs)s />

Remove Marker

Double click to zoom in and center on a location. Right click to set the marker on a position. You can also drag and drop the marker.

				""" % {
						'field_id': attrs['id'],
						'delete_id': 'map_delete_%s'%attrs['id'],
						'map_id': 'map_%s'%attrs['id'],
						'attrs': flatatt(final_attrs),
						'height': gmap_settings.HEIGHT,
						'width': gmap_settings.WIDTH,
						'zoom': gmap_settings.ZOOM,
						'center_lng': center_lng,
						'center_lat': center_lat,
					}
		return mark_safe(html)<|endoftext|>


<commit_before>#
# Kiwi: a Framework and Enhanced Widgets for Python
#
# Copyright (C) 2005 Async Open Source
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307
# USA
#
# Author(s): Johan Dahlin <jdahlin@async.com.br>
#

"""Distutils extensions and utilities"""

from distutils.command.install_data import install_data
from distutils.command.install_lib import install_lib
from distutils.dep_util import newer
from distutils.log import info, warn
from distutils.sysconfig import get_python_lib
from fnmatch import fnmatch
import os
import sys

class _VariableExtender:
    def __init__(self, distribution):
        install = distribution.get_command_obj('install')
        name = distribution.get_name()
        prefix = install.prefix
        self.datadir = os.path.join(prefix, 'share', name)
        if prefix == '/usr':
            self.sysconfdir = '/etc'
        else:
            self.sysconfdir = os.path.join(prefix, 'etc')

    def extend(self, string):
        for name, var in [('sysconfdir', self.sysconfdir),
                          ('datadir', self.datadir)]:
            string = string.replace('$' + name, var)
        return string

class KiwiInstallLib(install_lib):
    # Overridable by subclass
    resources = {}
    global_resources = {}

    def generate_template(self):
        if 'bdist_wininst' in sys.argv:
            prefix = 'import sys\nprefix = sys.prefix'
        else:
            install = self.distribution.get_command_obj('install')
            prefix = 'prefix = "%s"' % install.prefix

        name = self.distribution.get_name()
        filename = os.path.join(self.install_dir, name, '__installed__.py')
        self.mkpath(os.path.dirname(filename))
        fp = open(filename, 'w')
        fp.write('# Generated by setup.py do not modify\n')
        fp.write('import os\n')
        fp.write('%s\n' % prefix)
        self._write_dictionary(fp, 'resources', self.resources, prefix)
        self._write_dictionary(fp, 'global_resources', self.global_resources,
                               prefix)
        fp.close()

        return filename

    def _write_dictionary(self, fp, name, dictionary, prefix):
        name = self.distribution.get_name()
        fp.write('%s = {}\n' % name)
        for key, value in dictionary.items():
            value = self.varext.extend(value)
            parts = []
            for part in value.split('/'):
                if part == '$prefix':
                    part = 'prefix'
                else:
                    part = '"%s"' % part
                parts.append(part)

            fp.write("%s['%s'] = %s\n" % (
                name, key, 'os.path.join(%s)' % ', '.join(parts)))

    def install(self):
        self.varext = _VariableExtender(self.distribution)
        return install_lib.install(self) + [self.generate_template()]

# Backwards compat
TemplateInstallLib = KiwiInstallLib

class KiwiInstallData(install_data):
    def run(self):
        self.varext = _VariableExtender(self.distribution)

        # Extend variables in all data files
        data_files = []
        for target, files in self.data_files[:]:
            data_files.append((self.varext.extend(target), files))
        self.data_files = data_files
        return install_data.run(self)

def get_site_packages_dir(*dirs):
    """
    Gets the relative path of the site-packages directory

    This is mainly useful for setup.py usage:

    >>> setup(...
              data_files=[(get_site_packages_dir('foo'),
                           files..)])

    where files is a list of files to be installed in
    a directory called foo created in your site-packages directory

    @param dirs: directory names to be appended
    """

    libdir = get_python_lib(plat_specific=False,
                            standard_lib=True, prefix='')
    return os.path.join(libdir, 'site-packages', *dirs)

def listfiles(*dirs):
    """
    Lists all files in directories and optionally uses basic shell
    matching, example:

    >>> listfiles('data', 'glade', '*.glade')
    ['data/glade/Foo.glade', 'data/glade/Bar.glade', ...]

    @param dirs: directory parts
    """

    dir, pattern = os.path.split(os.path.join(*dirs))
    return [os.path.join(dir, filename)
            for filename in os.listdir(os.path.abspath(dir))
                if filename[0] != '.' and fnmatch(filename, pattern)]

def compile_po_files(domain, dirname='locale'):
    """
    Compiles po files to mo files.
    Note. this function depends on gettext utilities being installed

    @param domain: gettext domain
    @param dirname: base directory
    @returns: a list of po files
    """
    if os.system('msgfmt 2> /dev/null') != 256:
        warn('msgfmt is missing, not installing translations')
        return []

    data_files = []
    for po in listfiles('po', '*.po'):
        lang = os.path.basename(po[:-3])
        mo = os.path.join(dirname, lang, 'LC_MESSAGES', domain + '.mo')

        if not os.path.exists(mo) or newer(po, mo):
            directory = os.path.dirname(mo)
            if not os.path.exists(directory):
                info("creating %s" % directory)
                os.makedirs(directory)
            cmd = 'msgfmt -o %s %s' % (mo, po)
            info('compiling %s -> %s' % (po, mo))
            if os.system(cmd) != 0:
                raise SystemExit("Error while running msgfmt")
        dest = os.path.dirname(os.path.join('share', mo))
        data_files.append((dest, [mo]))

    return data_files

def listpackages(root, exclude=None):
    """Recursivly list all packages in directory root
    Optionally exclude can be specified which is a string
    like foo/bar.

    @param root: directory
    @param exclude: optional packages to be skipped
    """

    packages = []
    if not os.path.isdir(root):
        raise ValueError("root must be a directory")

    if os.path.exists(os.path.join(root, '__init__.py')):
        packages.append(root.replace('/', '.'))

    for filename in os.listdir(root):
        full = os.path.join(root, filename)
        if os.path.isdir(full):
            packages.extend(listpackages(full))

    if exclude:
        for package in packages[:]:
            if package.startswith(exclude):
                packages.remove(package)

    return packages

<commit_msg>Grr, wrong name<commit_after>#
# Kiwi: a Framework and Enhanced Widgets for Python
#
# Copyright (C) 2005 Async Open Source
#
# This library is free software; you can redistribute it and/or
# modify it under the terms of the GNU Lesser General Public
# License as published by the Free Software Foundation; either
# version 2.1 of the License, or (at your option) any later version.
#
# This library 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
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307
# USA
#
# Author(s): Johan Dahlin <jdahlin@async.com.br>
#

"""Distutils extensions and utilities"""

from distutils.command.install_data import install_data
from distutils.command.install_lib import install_lib
from distutils.dep_util import newer
from distutils.log import info, warn
from distutils.sysconfig import get_python_lib
from fnmatch import fnmatch
import os
import sys

class _VariableExtender:
    def __init__(self, distribution):
        install = distribution.get_command_obj('install')
        name = distribution.get_name()
        prefix = install.prefix
        self.datadir = os.path.join(prefix, 'share', name)
        if prefix == '/usr':
            self.sysconfdir = '/etc'
        else:
            self.sysconfdir = os.path.join(prefix, 'etc')

    def extend(self, string):
        for name, var in [('sysconfdir', self.sysconfdir),
                          ('datadir', self.datadir)]:
            string = string.replace('$' + name, var)
        return string

class KiwiInstallLib(install_lib):
    # Overridable by subclass
    resources = {}
    global_resources = {}

    def generate_template(self):
        if 'bdist_wininst' in sys.argv:
            prefix = 'import sys\nprefix = sys.prefix'
        else:
            install = self.distribution.get_command_obj('install')
            prefix = 'prefix = "%s"' % install.prefix

        name = self.distribution.get_name()
        filename = os.path.join(self.install_dir, name, '__installed__.py')
        self.mkpath(os.path.dirname(filename))
        fp = open(filename, 'w')
        fp.write('# Generated by setup.py do not modify\n')
        fp.write('import os\n')
        fp.write('%s\n' % prefix)
        self._write_dictionary(fp, 'resources', self.resources, prefix)
        self._write_dictionary(fp, 'global_resources', self.global_resources,
                               prefix)
        fp.close()

        return filename

    def _write_dictionary(self, fp, name, dictionary, prefix):
        fp.write('%s = {}\n' % name)
        for key, value in dictionary.items():
            value = self.varext.extend(value)
            parts = []
            for part in value.split('/'):
                if part == '$prefix':
                    part = 'prefix'
                else:
                    part = '"%s"' % part
                parts.append(part)

            fp.write("%s['%s'] = %s\n" % (
                name, key, 'os.path.join(%s)' % ', '.join(parts)))

    def install(self):
        self.varext = _VariableExtender(self.distribution)
        return install_lib.install(self) + [self.generate_template()]

# Backwards compat
TemplateInstallLib = KiwiInstallLib

class KiwiInstallData(install_data):
    def run(self):
        self.varext = _VariableExtender(self.distribution)

        # Extend variables in all data files
        data_files = []
        for target, files in self.data_files[:]:
            data_files.append((self.varext.extend(target), files))
        self.data_files = data_files
        return install_data.run(self)

def get_site_packages_dir(*dirs):
    """
    Gets the relative path of the site-packages directory

    This is mainly useful for setup.py usage:

    >>> setup(...
              data_files=[(get_site_packages_dir('foo'),
                           files..)])

    where files is a list of files to be installed in
    a directory called foo created in your site-packages directory

    @param dirs: directory names to be appended
    """

    libdir = get_python_lib(plat_specific=False,
                            standard_lib=True, prefix='')
    return os.path.join(libdir, 'site-packages', *dirs)

def listfiles(*dirs):
    """
    Lists all files in directories and optionally uses basic shell
    matching, example:

    >>> listfiles('data', 'glade', '*.glade')
    ['data/glade/Foo.glade', 'data/glade/Bar.glade', ...]

    @param dirs: directory parts
    """

    dir, pattern = os.path.split(os.path.join(*dirs))
    return [os.path.join(dir, filename)
            for filename in os.listdir(os.path.abspath(dir))
                if filename[0] != '.' and fnmatch(filename, pattern)]

def compile_po_files(domain, dirname='locale'):
    """
    Compiles po files to mo files.
    Note. this function depends on gettext utilities being installed

    @param domain: gettext domain
    @param dirname: base directory
    @returns: a list of po files
    """
    if os.system('msgfmt 2> /dev/null') != 256:
        warn('msgfmt is missing, not installing translations')
        return []

    data_files = []
    for po in listfiles('po', '*.po'):
        lang = os.path.basename(po[:-3])
        mo = os.path.join(dirname, lang, 'LC_MESSAGES', domain + '.mo')

        if not os.path.exists(mo) or newer(po, mo):
            directory = os.path.dirname(mo)
            if not os.path.exists(directory):
                info("creating %s" % directory)
                os.makedirs(directory)
            cmd = 'msgfmt -o %s %s' % (mo, po)
            info('compiling %s -> %s' % (po, mo))
            if os.system(cmd) != 0:
                raise SystemExit("Error while running msgfmt")
        dest = os.path.dirname(os.path.join('share', mo))
        data_files.append((dest, [mo]))

    return data_files

def listpackages(root, exclude=None):
    """Recursivly list all packages in directory root
    Optionally exclude can be specified which is a string
    like foo/bar.

    @param root: directory
    @param exclude: optional packages to be skipped
    """

    packages = []
    if not os.path.isdir(root):
        raise ValueError("root must be a directory")

    if os.path.exists(os.path.join(root, '__init__.py')):
        packages.append(root.replace('/', '.'))

    for filename in os.listdir(root):
        full = os.path.join(root, filename)
        if os.path.isdir(full):
            packages.extend(listpackages(full))

    if exclude:
        for package in packages[:]:
            if package.startswith(exclude):
                packages.remove(package)

    return packages

<|endoftext|>


<commit_before>'''
Pushes stored user favorites to Scout for alert tracking.
'''
import json
import logging
import urlparse
from optparse import make_option

from django.contrib.auth.models import User
from django.core.management.base import NoArgsCommand

from billy.core import settings, user_db
from billy.utils import JSONEncoderPlus

import requests

_log = logging.getLogger('billy.web.public.management.commands.scout_push')

class Command(NoArgsCommand):
    option_list = NoArgsCommand.option_list + (
        make_option('--dry-run', action='store_true', dest='dry_run',
                    default=False, help='dry run'),
    )

    def handle_noargs(self, **options):
        # TODO: this would be nicer if it were by email address
        usernames = user_db.favorites.distinct('username')

        for user in usernames:
            self.push_user(username, options['dry_run'])

    def push_user(self, username, dry_run):
        user = User.objects.get(username=username)
        payload = {'email': user.email,
                   'secret_key': settings.SCOUT_SECRET_KEY,
                   'service': settings.SCOUT_SERVICE,
                   'notifications': 'email_daily'}
        interests = []
        for favorite in user_db.favorites.find({'username': username}):
            if favorite['obj_type'] in ('legislator', 'bill', 'committee'):
                interest = {'interest_type': 'item',
                            'item_type': 'state_' + favorite['obj_type'],
                            'item_id': favorite['obj_id'],
                           }
            elif favorite['obj_type'] == 'search':
                params = urlparse.parse_qs(favorite['search_params'])
                interest = {
                    'interest_type': 'search',
                    'search_type': 'state_bills',
                    'query_type': 'advanced',
                    'filters': self._translate_filter_data(favorite, params)
                }
                if 'search_text' in params:
                    interest['in'] = params.get('search_text').pop()
                else:
                    interest['in'] = ''

            else:
                _log.warning('Unknown favorite type: %s', favorite['obj_type'])
                continue

            interest['active'] = favorite['is_favorite']
            interest['changed_at'] = favorite['timestamp']
            interests.append(interest)

        payload['interests'] = interests
        _log.info('pushing %s interests for %s', len(interests),
                  payload['email'])

        if not dry_run:
            url = 'http://scout.sunlightfoundation.com/remote/service/sync'
            payload = json.dumps(payload, cls=JSONEncoderPlus)
            requests.post(url, data=payload)

    def _translate_filter_data(self, favorite, params):
        '''Edit the favorite['search_params'] object and make them
        match the param names used in an api request.
        '''
        # some api params have no analog in the front-end search: updated_since
        api_param_name_set = set([
            'state',
            'search_window',
            'chamber',
            'subjects',
            'sponsor_id',
            'session',
            'type',
            'status'])

        result = {}

        # Rename certain front-end parameters to their api equivalents.
        api_param_names = {
            'session': 'search_window',
            'sponsor__leg_id': 'sponsor_id'
        }

        for k, v in params.items():

            if k == 'session':
                v = 'session:' + v.pop()

            api_param_name = api_param_names.get(k, k)

            if api_param_name in api_param_name_set:

                # Flatten any single-item param arrays into strings.
                if isinstance(v, list) and k not in ['status', 'subjects']:
                    if len(v) == 1:
                        v = v.pop()

                result[api_param_name] = v

        # Add the state abbreviation.
        if 'search_abbr' in favorite:
            result['state'] = favorite['search_abbr'].upper()

        return result

<commit_msg>fix billy<commit_after>'''
Pushes stored user favorites to Scout for alert tracking.
'''
import json
import logging
import urlparse
from optparse import make_option

from django.contrib.auth.models import User
from django.core.management.base import NoArgsCommand

from billy.core import settings, user_db
from billy.utils import JSONEncoderPlus

import requests

_log = logging.getLogger('billy.web.public.management.commands.scout_push')

class Command(NoArgsCommand):
    option_list = NoArgsCommand.option_list + (
        make_option('--dry-run', action='store_true', dest='dry_run',
                    default=False, help='dry run'),
    )

    def handle_noargs(self, **options):
        # TODO: this would be nicer if it were by email address
        usernames = user_db.favorites.distinct('username')

        for username in usernames:
            self.push_user(username, options['dry_run'])

    def push_user(self, username, dry_run):
        user = User.objects.get(username=username)
        payload = {'email': user.email,
                   'secret_key': settings.SCOUT_SECRET_KEY,
                   'service': settings.SCOUT_SERVICE,
                   'notifications': 'email_daily'}
        interests = []
        for favorite in user_db.favorites.find({'username': username}):
            if favorite['obj_type'] in ('legislator', 'bill', 'committee'):
                interest = {'interest_type': 'item',
                            'item_type': 'state_' + favorite['obj_type'],
                            'item_id': favorite['obj_id'],
                           }
            elif favorite['obj_type'] == 'search':
                params = urlparse.parse_qs(favorite['search_params'])
                interest = {
                    'interest_type': 'search',
                    'search_type': 'state_bills',
                    'query_type': 'advanced',
                    'filters': self._translate_filter_data(favorite, params)
                }
                if 'search_text' in params:
                    interest['in'] = params.get('search_text').pop()
                else:
                    interest['in'] = ''

            else:
                _log.warning('Unknown favorite type: %s', favorite['obj_type'])
                continue

            interest['active'] = favorite['is_favorite']
            interest['changed_at'] = favorite['timestamp']
            interests.append(interest)

        payload['interests'] = interests
        _log.info('pushing %s interests for %s', len(interests),
                  payload['email'])

        if not dry_run:
            url = 'http://scout.sunlightfoundation.com/remote/service/sync'
            payload = json.dumps(payload, cls=JSONEncoderPlus)
            requests.post(url, data=payload)

    def _translate_filter_data(self, favorite, params):
        '''Edit the favorite['search_params'] object and make them
        match the param names used in an api request.
        '''
        # some api params have no analog in the front-end search: updated_since
        api_param_name_set = set([
            'state',
            'search_window',
            'chamber',
            'subjects',
            'sponsor_id',
            'session',
            'type',
            'status'])

        result = {}

        # Rename certain front-end parameters to their api equivalents.
        api_param_names = {
            'session': 'search_window',
            'sponsor__leg_id': 'sponsor_id'
        }

        for k, v in params.items():

            if k == 'session':
                v = 'session:' + v.pop()

            api_param_name = api_param_names.get(k, k)

            if api_param_name in api_param_name_set:

                # Flatten any single-item param arrays into strings.
                if isinstance(v, list) and k not in ['status', 'subjects']:
                    if len(v) == 1:
                        v = v.pop()

                result[api_param_name] = v

        # Add the state abbreviation.
        if 'search_abbr' in favorite:
            result['state'] = favorite['search_abbr'].upper()

        return result

<|endoftext|>


<commit_before># -*- encoding: utf-8 -*-
from __future__ import unicode_literals

from django.contrib import messages
from django.core.urlresolvers import reverse
from django.db import transaction
from django.db.models import get_model
from django.http import HttpResponseRedirect
from django.views.generic import (
    CreateView,
    DeleteView,
    UpdateView,
    TemplateView,
)

from braces.views import (
    LoginRequiredMixin,
    StaffuserRequiredMixin,
)

from base.view_utils import BaseMixin

from .models import (
    BlockError,
    Page,
    PageSection,
    Section,
)

class ContentPageMixin(BaseMixin):
    """Page information."""

    def get_context_data(self, **kwargs):
        context = super(ContentPageMixin, self).get_context_data(**kwargs)
        context.update(dict(
            page=self.get_page(),
            pages=Page.objects.menu(),
        ))
        return context

    def get_page(self):
        menu = self.kwargs.get('menu', None)
        page = self.kwargs.get('page', None)
        if not page:
            raise BlockError("no 'page' parameter in url")
        try:
            if menu:
                return Page.objects.get(slug=page, slug_menu=menu)
            else:
                return Page.objects.get(slug=page)
        except Page.DoesNotExist:
            raise BlockError(
                "Page '{}' (menu '{}') does not exist".format(page, menu)
            )

    def get_page_section(self):
        page = self.get_page()
        section = self.get_section()
        try:
            return PageSection.objects.get(page=page, section=section)
        except PageSection.DoesNotExist:
            raise BlockError(
                "Page section '{}/{}' does not exist".format(
                    page.slug,
                    section.slug,
                )
            )

    def get_section(self):
        section = self.kwargs.get('section', None)
        if not section:
            raise BlockError("no 'section' parameter in url")
        try:
            return Section.objects.get(slug=section)
        except Section.DoesNotExist:
            raise BlockError("Section '{}' does not exist".format(section))

    def get_success_url(self):
        return self.object.block.page_section.page.get_absolute_url()

class ContentCreateView(ContentPageMixin, BaseMixin, CreateView):

    def form_valid(self, form):
        self.object = form.save(commit=False)
        try:
            block = self.block_class(page_section=self.get_page_section())
        except AttributeError:
            raise BlockError(
                "You need to add the 'block_class' attribute "
                "to your '{}' class.".format(self.__class__.__name__)
            )
        block.save()
        self.object.block = block
        self.object.order = self.model.objects.next_order()
        return super(ContentCreateView, self).form_valid(form)

class ContentPublishView(BaseMixin, UpdateView):

    def form_valid(self, form):
        """Publish 'pending' content."""
        self.object = form.save(commit=False)
        self.object.block.publish(self.request.user)
        messages.info(
            self.request,
            "Published block, id {}".format(
                self.object.block.pk,
            )
        )
        return HttpResponseRedirect(self.get_success_url())

    def get_success_url(self):
        return self.object.block.page_section.page.get_absolute_url()

class ContentRemoveView(BaseMixin, UpdateView):

    def form_valid(self, form):
        self.object = form.save(commit=False)
        self.object.block.remove(self.request.user)
        messages.info(
            self.request,
            "Removed block {}".format(
                self.object.block.pk,
            )
        )
        return HttpResponseRedirect(self.get_success_url())

    def get_success_url(self):
        return self.object.block.page_section.page.get_absolute_url()

class ContentUpdateView(BaseMixin, UpdateView):

    def form_valid(self, form):
        self.object = form.save(commit=False)
        self.object.set_pending_edit()
        return super(ContentUpdateView, self).form_valid(form)

    def get_context_data(self, **kwargs):
        context = super(ContentUpdateView, self).get_context_data(**kwargs)
        context.update(dict(
            is_update=True,
        ))
        return context

    def get_section(self):
        return self.object.block.section

    def get_success_url(self):
        return self.object.block.page_section.page.get_absolute_url()

class ElementCreateView(BaseMixin, CreateView):
    """Update the parent (content) object so it knows it has been changed."""

    def form_valid(self, form):
        self.object = form.save(commit=False)
        content = self.object.get_parent()
        content.set_pending_edit()
        content.save()
        return super(ElementCreateView, self).form_valid(form)

class ElementDeleteView(BaseMixin, DeleteView):
    """Update the parent (content) object so it knows it has been changed."""

    def delete(self, request, *args, **kwargs):
        with transaction.atomic():
            result = super(ElementDeleteView, self).delete(
                request, *args, **kwargs
            )
            content = self.object.get_parent()
            content.set_pending_edit()
            content.save()
        return result

class ElementUpdateView(BaseMixin, UpdateView):
    """Update the parent (content) object so it knows it has been changed."""

    def form_valid(self, form):
        self.object = form.save(commit=False)
        content = self.object.get_parent()
        content.set_pending_edit()
        content.save()
        return super(ElementUpdateView, self).form_valid(form)

class PageTemplateMixin(object):

    def get_template_names(self) :
        page = self.get_page()
        if page.template_name:
            template_name = page.template_name
        else:
            raise BlockError(
                "Page '{}' has no 'template_name'".format(page.slug)
            )
        return [template_name,]

class PageDesignMixin(object):

    def get_context_data(self, **kwargs):
        context = super(PageDesignMixin, self).get_context_data(**kwargs)
        page = self.get_page()
        context.update(dict(design=True))
        for e in PageSection.objects.filter(page=page) :
            block_create_url = '{}_create_url'.format(e.section.slug)
            block_list_name = '{}_list'.format(e.section.slug)
            block_model = get_model(e.block_app, e.block_model)
            kwargs = dict(section=e.section.slug)
            kwargs.update(page.get_url_kwargs())
            context.update({
                block_list_name: block_model.objects.pending(e),
            })
            if e.create_url_name:
                context.update({
                    block_create_url: reverse(e.create_url_name, kwargs=kwargs),
                })
        return context

class PageDesignView(
        LoginRequiredMixin, StaffuserRequiredMixin,
        PageDesignMixin, PageTemplateMixin, ContentPageMixin, TemplateView):

    pass

class PageMixin(object):

    def get_context_data(self, **kwargs):
        context = super(PageMixin, self).get_context_data(**kwargs)
        page = self.get_page()
        context.update(dict(design=False))
        for e in PageSection.objects.filter(page=page):
            block_list_name = '{}_list'.format(e.section.slug)
            block_model = get_model(e.block_app, e.block_model)
            context.update({
                block_list_name: block_model.objects.published(e),
            })
        return context

class PageView(
        PageMixin, PageTemplateMixin, ContentPageMixin,
        BaseMixin, TemplateView):

    pass
<commit_msg>duplicate class mixin.  tidy message.<commit_after># -*- encoding: utf-8 -*-
from __future__ import unicode_literals

from django.contrib import messages
from django.core.urlresolvers import reverse
from django.db import transaction
from django.db.models import get_model
from django.http import HttpResponseRedirect
from django.views.generic import (
    CreateView,
    DeleteView,
    UpdateView,
    TemplateView,
)

from braces.views import (
    LoginRequiredMixin,
    StaffuserRequiredMixin,
)

from base.view_utils import BaseMixin

from .models import (
    BlockError,
    Page,
    PageSection,
    Section,
)

class ContentPageMixin(BaseMixin):
    """Page information."""

    def get_context_data(self, **kwargs):
        context = super(ContentPageMixin, self).get_context_data(**kwargs)
        context.update(dict(
            page=self.get_page(),
            pages=Page.objects.menu(),
        ))
        return context

    def get_page(self):
        menu = self.kwargs.get('menu', None)
        page = self.kwargs.get('page', None)
        if not page:
            raise BlockError("no 'page' parameter in url")
        try:
            if menu:
                return Page.objects.get(slug=page, slug_menu=menu)
            else:
                return Page.objects.get(slug=page)
        except Page.DoesNotExist:
            raise BlockError(
                "Page '{}', menu '{}' does not exist".format(page, menu)
            )

    def get_page_section(self):
        page = self.get_page()
        section = self.get_section()
        try:
            return PageSection.objects.get(page=page, section=section)
        except PageSection.DoesNotExist:
            raise BlockError(
                "Page section '{}/{}' does not exist".format(
                    page.slug,
                    section.slug,
                )
            )

    def get_section(self):
        section = self.kwargs.get('section', None)
        if not section:
            raise BlockError("no 'section' parameter in url")
        try:
            return Section.objects.get(slug=section)
        except Section.DoesNotExist:
            raise BlockError("Section '{}' does not exist".format(section))

    def get_success_url(self):
        return self.object.block.page_section.page.get_absolute_url()

class ContentCreateView(ContentPageMixin, BaseMixin, CreateView):

    def form_valid(self, form):
        self.object = form.save(commit=False)
        try:
            block = self.block_class(page_section=self.get_page_section())
        except AttributeError:
            raise BlockError(
                "You need to add the 'block_class' attribute "
                "to your '{}' class.".format(self.__class__.__name__)
            )
        block.save()
        self.object.block = block
        self.object.order = self.model.objects.next_order()
        return super(ContentCreateView, self).form_valid(form)

class ContentPublishView(BaseMixin, UpdateView):

    def form_valid(self, form):
        """Publish 'pending' content."""
        self.object = form.save(commit=False)
        self.object.block.publish(self.request.user)
        messages.info(
            self.request,
            "Published block, id {}".format(
                self.object.block.pk,
            )
        )
        return HttpResponseRedirect(self.get_success_url())

    def get_success_url(self):
        return self.object.block.page_section.page.get_absolute_url()

class ContentRemoveView(BaseMixin, UpdateView):

    def form_valid(self, form):
        self.object = form.save(commit=False)
        self.object.block.remove(self.request.user)
        messages.info(
            self.request,
            "Removed block {}".format(
                self.object.block.pk,
            )
        )
        return HttpResponseRedirect(self.get_success_url())

    def get_success_url(self):
        return self.object.block.page_section.page.get_absolute_url()

class ContentUpdateView(BaseMixin, UpdateView):

    def form_valid(self, form):
        self.object = form.save(commit=False)
        self.object.set_pending_edit()
        return super(ContentUpdateView, self).form_valid(form)

    def get_context_data(self, **kwargs):
        context = super(ContentUpdateView, self).get_context_data(**kwargs)
        context.update(dict(
            is_update=True,
        ))
        return context

    def get_section(self):
        return self.object.block.section

    def get_success_url(self):
        return self.object.block.page_section.page.get_absolute_url()

class ElementCreateView(BaseMixin, CreateView):
    """Update the parent (content) object so it knows it has been changed."""

    def form_valid(self, form):
        self.object = form.save(commit=False)
        content = self.object.get_parent()
        content.set_pending_edit()
        content.save()
        return super(ElementCreateView, self).form_valid(form)

class ElementDeleteView(BaseMixin, DeleteView):
    """Update the parent (content) object so it knows it has been changed."""

    def delete(self, request, *args, **kwargs):
        with transaction.atomic():
            result = super(ElementDeleteView, self).delete(
                request, *args, **kwargs
            )
            content = self.object.get_parent()
            content.set_pending_edit()
            content.save()
        return result

class ElementUpdateView(BaseMixin, UpdateView):
    """Update the parent (content) object so it knows it has been changed."""

    def form_valid(self, form):
        self.object = form.save(commit=False)
        content = self.object.get_parent()
        content.set_pending_edit()
        content.save()
        return super(ElementUpdateView, self).form_valid(form)

class PageTemplateMixin(object):

    def get_template_names(self) :
        page = self.get_page()
        if page.template_name:
            template_name = page.template_name
        else:
            raise BlockError(
                "Page '{}' has no 'template_name'".format(page.slug)
            )
        return [template_name,]

class PageDesignMixin(object):

    def get_context_data(self, **kwargs):
        context = super(PageDesignMixin, self).get_context_data(**kwargs)
        page = self.get_page()
        context.update(dict(design=True))
        for e in PageSection.objects.filter(page=page) :
            block_create_url = '{}_create_url'.format(e.section.slug)
            block_list_name = '{}_list'.format(e.section.slug)
            block_model = get_model(e.block_app, e.block_model)
            kwargs = dict(section=e.section.slug)
            kwargs.update(page.get_url_kwargs())
            context.update({
                block_list_name: block_model.objects.pending(e),
            })
            if e.create_url_name:
                context.update({
                    block_create_url: reverse(e.create_url_name, kwargs=kwargs),
                })
        return context

class PageDesignView(
        LoginRequiredMixin, StaffuserRequiredMixin,
        PageDesignMixin, PageTemplateMixin, ContentPageMixin, TemplateView):

    pass

class PageMixin(object):

    def get_context_data(self, **kwargs):
        context = super(PageMixin, self).get_context_data(**kwargs)
        page = self.get_page()
        context.update(dict(design=False))
        for e in PageSection.objects.filter(page=page):
            block_list_name = '{}_list'.format(e.section.slug)
            block_model = get_model(e.block_app, e.block_model)
            context.update({
                block_list_name: block_model.objects.published(e),
            })
        return context

class PageView(
        PageMixin, PageTemplateMixin, ContentPageMixin, TemplateView):

    pass
<|endoftext|>