<commit_before>import operator as op
from decimal import Decimal, getcontext
getcontext().prec = 6

from .globals import BLK_OPEN, BLK_CLOSE, COMMENT_OPEN, COMMENT_CLOSE
from .types import (RuleBlock, CSSRule, MixIn, UnpackMe, RootBlock,
                    Distance, ParentReference, Color, AtRuleBlock, Keyframe)
from .core import GDef
from .functions import builtin_funcs
from .errors import UnknownMixinOrFunc

@GDef(r'ltree = root_element+')
def ltree(env, node, children):
    return children

@GDef(r'root_element = ___ (root_block / at_rule / mixin_decl / var_decl / rule_block) ___')
def root_element(env, node, children):
    return children[1][0]

### SELECTORS ###

@GDef(r'rule_block = tag? simple_selector ("," _ simple_selector)* ___ block')
def rule_block(env, node, children):
    tag, simple_sel, more_selectors, _, block = children
    tag = tag[0] if tag else None

    selectors = [simple_sel]
    if more_selectors:
        selectors = selectors.extend([selector for _, selector in more_selectors])

    return RuleBlock(tag=tag,
                     selectors=selectors,
                     block=block,
                     index=node.start)

@GDef(r'block = blk_open ___ block_element+ blk_close')
def block(env, node, children):
    return children[2]

@GDef(r'block_element = (mixin_or_func_call / declaration / rule_block / parent_reference) ___')
def block_element(env, node, children):
    return children[0][0]

@GDef(r'simple_selector = ((type_sel / universal_sel) (attribute_sel / id_class_sel / pseudo_class)*) / (attribute_sel / id_class_sel)')
def simple_selector(env, node, children):
    return node.text

@GDef(r'''
type_sel = name
universal_sel = "*"

attribute_sel = "[" name ("=" / "~=" / "|=") name "]"
id_class_sel = ("#" / ".") name

pseudo_class = ":" ((pseudo_class_param "(" num ")") / pseudo_class_noparam)
pseudo_class_param = "nth-child" / "nth-last-child" / "nth-of-type" / "nth-last-of-type" / "lang"
pseudo_class_noparam = "last-child" / "first-of-type" / "last-of-type" / "only-child" / "only-of-type" / "root" / "empty" / "target" / "enabled" / "disabled" / "checked" / "link" / "visited" / "hover" / "active" / "focus" / "first-letter" / "first-line" / "first-child" / "before" / "after"
''')
def selector_misc(env, node, children):
    return node.text

### PARENT REFERENCE ###

@GDef(r'parent_reference = "&" parent_block')
def parent_reference(env, node, children):
    return ParentReference(children[1])

@GDef(r'parent_block = tag? (simple_selector / pseudo_class) ("," _ (simple_selector / pseudo_class))* ___ block')
def parent_block(env, node, children):
    tag, (simple_sel, ), more_selectors, _, block = children
    tag = tag[0] if tag else None

    selectors = [simple_sel]
    if more_selectors:
        selectors = selectors.extend([selector for _, (selector, ) in more_selectors])

    return RuleBlock(tag=tag,
                     selectors=selectors,
                     block=block,
                     index=node.start)

### AT-RULES ###

@GDef(r'at_rule = tag? "@" (at_rule_keyframe / at_rule_normal)')
def at_rule(env, node, children):
    tag = children[0][0] if children[0] else None
    text, block = children[2][0]

    return AtRuleBlock(tag=tag,
                       text=text.strip(),
                       block=block,
                       index=node.start)

@GDef(r'at_rule_normal = any ___ block')
def at_rule_normal(env, node, children):
    text, _, block = children
    return (text, block)

@GDef(r'at_rule_keyframe = "keyframes" _ name ___ keyframes_block')
def at_rule_keyframes(env, node, children):
    _, _, name, ___, block = children
    return ("keyframes " + name, block)

@GDef(r'keyframes_block = blk_open ___ keyframe+ blk_close')
def keyframes_block(env, node, children):
    return children[2]

@GDef(r'keyframe = tag? expr ___ block ___')
def keyframe(env, node, children):
    tag, condition, _, block, _ = children
    tag = tag[0] if tag else None
    return Keyframe(
        tag=tag,
        condition=condition,
        block=block,
        index=node.start)

### CSS RULES ###

@GDef(r'declaration = tag? property ":" _ expr+')
def declaration(env, node, children):
    tag, prop, _, _, values = children
    tag = tag[0] if tag else None
    return CSSRule(tag=tag,
                   prop=prop,
                   values=values,
                   index=node.start)

@GDef(r'property = ~"[a-zA-Z\_][a-zA-Z0-9\-\_]*"')
def property_(env, node, children):
    return node.text

@GDef(r'expr = (mixin_or_func_call / math / color / string_val) _?')
def expr(env, node, children):
    return children[0][0]

### MIXINS, FUNCTIONS, and VARIABLES ###

@GDef(r'mixin_decl = name "(" (_? name _?)* "):" ___ block',
      defer=True)
def mixin_decl(env, node):
    name, _, variables, _, _, block = node

    ly_engine = LY(env=env)
    name = ly_engine._evalnode(name)
    variables = [varname for _, varname, _ in ly_engine._evalnode(variables)]

    def f(*args):
        global_vars = list(env.items())
        local_vars = list(zip(variables, args))
        temp_env = dict(global_vars + local_vars)

        ly_engine_local = LY(env=temp_env)
        return ly_engine_local._evalnode(block)

    env[name] = MixIn(name=name, func=f)

@GDef(r'mixin_or_func_call = name "(" (_? expr _?)* ")"')
def mixin_or_func_call(env, node, children):
    name, _, args, _ = children
    args = [arg for _, arg, _ in args]

    if name in builtin_funcs:
        return builtin_funcs[name](*args)
    elif name in env:
        return UnpackMe(env[name](*args))

    raise UnknownMixinOrFunc(location=node.start)

@GDef(r'lvalue = ~"[a-zA-Z\_][a-zA-Z0-9\-\_]*"')
def lvalue(env, node, children):
    if node.text in env:
        return env[node.text]
    return node.text

@GDef(r'var_decl = name _? "=" _? expr ___')
def var_decl(env, node, children):
    name, _, _, _, value, _ = children
    env[name] = value

### ROOT BLOCKS ###

@GDef(r'root_block = "(root" ("." name)? ")" _? string_val?')
def root_block(env, node, children):
    _, possible_name, _, _, prefix = children
    tag_name = possible_name[0][1] if possible_name else None
    if not prefix:
        prefix = ''
    else:
        prefix = prefix[0]
    return RootBlock(tag_name=tag_name, prefix=prefix)

### TEXT ###

@GDef(r'''string_val = ~'\"[^"\n]*\"' ''')
def string_val(env, node, children):
    return node.text[1:-1]

### NUMERIC ###

@GDef(r'math = sum')
def math(env, node, children):
    return children[0]

@GDef(r'''
sum = prod (_? sum_op _? prod)*
prod = equality (_ prod_op _ equality)*
equality = value (_ equality_op _ value)*
''')
def math_operation(env, node, children):
    if len(children[1]):
        operations = [
            (operator, operand) for _, operator, _, operand in children[1]]
        return do_math(children[0], operations)
    return children[0]

@GDef(r'''
sum_op = "+" / "-"
prod_op = "*" / "/"
equality_op = "=="
''')
def operator_symbols(env, node, children):
    return node.text

@GDef(r'value = num_val / lvalue / paren')
def math_value(env, node, children):
    return children[0]

@GDef(r'paren = "(" _? sum _? ")"')
def math_paren(env, node, children):
    return children[2]

@GDef(r'num_val = distance / num / hexcolor / color_name / lvalue')
def num_val(env, node, children):
    return children[0]

@GDef(r'distance = num unit')
def distance(env, node, children):
    num, unit = children
    return Distance(value=num, unit=unit)

@GDef(r'num = ~"\-?\d+(\.\d+)?"')
def num(env, node, children):
    return Decimal(node.text)

@GDef(r'unit = "em" / "ex" / "px" / "cm" / "mm" / "in" / "pt" / "pc" / "deg" / "rad" / "grad" / "ms" / "s" / "hz" / "khz" / "%"')
def unit(env, node, children):
    return node.text

### MATH HELPER FUNCTION ###

OPERATORS = {'+': op.add,
             '-': op.sub,
             '*': op.mul,
             '/': op.truediv,
             '==': op.eq
             }

def do_math(start, operations):
    output_val = start
    for optr_symbol, operand in operations:
        output_val = OPERATORS[optr_symbol](output_val, operand)
    return output_val

### COLORS ###

@GDef(r'color = color_name / hexcolor')
def color(env, node, children):
    return node.text

@GDef(r'hexcolor = "#" ~"[0-9a-fA-F]+"')
def color_(env, node, children):
    return Color(node.text, ctype='hex')

@GDef(r'color_name = "aliceblue" / "antiquewhite" / "aqua" / "aquamarine" / "azure" / "beige" / "bisque" / "black" / "blanchedalmond" / "blue" / "blueviolet" / "brown" / "burlywood" / "cadetblue" / "chartreuse" / "chocolate" / "coral" / "cornflowerblue" / "cornsilk" / "crimson" / "cyan" / "darkblue" / "darkcyan" / "darkgoldenrod" / "darkgray" / "darkgreen" / "darkkhaki" / "darkmagenta" / "darkolivegreen" / "darkorange" / "darkorchid" / "darkred" / "darksalmon" / "darkseagreen" / "darkslateblue" / "darkslategray" / "darkturquoise" / "darkviolet" / "deeppink" / "deepskyblue" / "dimgray" / "dodgerblue" / "firebrick" / "floralwhite" / "forestgreen" / "fuchsia" / "gainsboro" / "ghostwhite" / "gold" / "goldenrod" / "gray" / "green" / "greenyellow" / "honeydew" / "hotpink" / "indianred " / "indigo " / "ivory" / "khaki" / "lavender" / "lavenderblush" / "lawngreen" / "lemonchiffon" / "lightblue" / "lightcoral" / "lightcyan" / "lightgoldenrodyellow" / "lightgray" / "lightgreen" / "lightpink" / "lightsalmon" / "lightseagreen" / "lightskyblue" / "lightslategray" / "lightsteelblue" / "lightyellow" / "lime" / "limegreen" / "linen" / "magenta" / "maroon" / "mediumaquamarine" / "mediumblue" / "mediumorchid" / "mediumpurple" / "mediumseagreen" / "mediumslateblue" / "mediumspringgreen" / "mediumturquoise" / "mediumvioletred" / "midnightblue" / "mintcream" / "mistyrose" / "moccasin" / "navajowhite" / "navy" / "oldlace" / "olive" / "olivedrab" / "orange" / "orangered" / "orchid" / "palegoldenrod" / "palegreen" / "paleturquoise" / "palevioletred" / "papayawhip" / "peachpuff" / "peru" / "pink" / "plum" / "powderblue" / "purple" / "red" / "rosybrown" / "royalblue" / "saddlebrown" / "salmon" / "sandybrown" / "seagreen" / "seashell" / "sienna" / "silver" / "skyblue" / "slateblue" / "slategray" / "snow" / "springgreen" / "steelblue" / "tan" / "teal" / "thistle" / "tomato" / "turquoise" / "violet" / "wheat" / "white" / "whitesmoke" / "yellow" / "yellowgreen"')
def color_name(env, node, children):
    return Color(node.text, ctype='named')

### BASIC ###

@GDef(r'tag = "(" name ")" _')
def tag(env, node, children):
    return children[1]

@GDef(r'name = ~"[a-zA-Z\_][a-zA-Z0-9\-\_]*"')
def name(env, node, children):
    return node.text

@GDef(r'any = ~".+"')
def any(env, node, children):
    return node.text.strip()

@GDef(r'''
___ = ~"[\s]*" (comment ~"[\s]*")*
_ = ~"[ \t]+"
blk_open = "{blk_open}"
blk_close = "{blk_close}"
comment = ~"{comment_open}.*(?!{comment_close})"
'''.format(
    comment_open=COMMENT_OPEN,
    comment_close=COMMENT_CLOSE,
    blk_open=BLK_OPEN,
    blk_close=BLK_CLOSE)
)
def return_none(env, node, children):
    return None

from .core import LY
<commit_msg>allow child selector, immediate sibling selector, general sibling selector<commit_after>import operator as op
from decimal import Decimal, getcontext
getcontext().prec = 6

from .globals import BLK_OPEN, BLK_CLOSE, COMMENT_OPEN, COMMENT_CLOSE
from .types import (RuleBlock, CSSRule, MixIn, UnpackMe, RootBlock,
                    Distance, ParentReference, Color, AtRuleBlock, Keyframe)
from .core import GDef
from .functions import builtin_funcs
from .errors import UnknownMixinOrFunc

@GDef(r'ltree = root_element+')
def ltree(env, node, children):
    return children

@GDef(r'root_element = ___ (root_block / at_rule / mixin_decl / var_decl / rule_block) ___')
def root_element(env, node, children):
    return children[1][0]

### SELECTORS ###

@GDef(r'rule_block = tag? simple_selector ("," _ simple_selector)* ___ block')
def rule_block(env, node, children):
    tag, simple_sel, more_selectors, _, block = children
    tag = tag[0] if tag else None

    selectors = [simple_sel]
    if more_selectors:
        selectors = selectors.extend([selector for _, selector in more_selectors])

    return RuleBlock(tag=tag,
                     selectors=selectors,
                     block=block,
                     index=node.start)

@GDef(r'block = blk_open ___ block_element+ blk_close')
def block(env, node, children):
    return children[2]

@GDef(r'block_element = (mixin_or_func_call / declaration / rule_block / parent_reference) ___')
def block_element(env, node, children):
    return children[0][0]

@GDef(r'simple_selector = selector_with_space? (((type_sel / universal_sel) (attribute_sel / id_class_sel / pseudo_class)*) / (attribute_sel / id_class_sel))')
def simple_selector(env, node, children):
    return node.text

@GDef(r'''
type_sel = name
universal_sel = "*"

attribute_sel = "[" name ("=" / "~=" / "|=") name "]"
id_class_sel = ("#" / ".") name

pseudo_class = ":" ((pseudo_class_param "(" num ")") / pseudo_class_noparam)
pseudo_class_param = "nth-child" / "nth-last-child" / "nth-of-type" / "nth-last-of-type" / "lang"
pseudo_class_noparam = "last-child" / "first-of-type" / "last-of-type" / "only-child" / "only-of-type" / "root" / "empty" / "target" / "enabled" / "disabled" / "checked" / "link" / "visited" / "hover" / "active" / "focus" / "first-letter" / "first-line" / "first-child" / "before" / "after"
''')
def selector_misc(env, node, children):
    return node.text

@GDef(r'selector_with_space = _? (">" / "+" / "~") _?')
def selector_with_space(env, node, children):
    'Support immediate child, immediate sibling, and general sibling selectors.'
    return node.text.strip()

### PARENT REFERENCE ###

@GDef(r'parent_reference = "&" parent_block')
def parent_reference(env, node, children):
    return ParentReference(children[1])

@GDef(r'parent_block = tag? (simple_selector / pseudo_class) ("," _ (simple_selector / pseudo_class))* ___ block')
def parent_block(env, node, children):
    tag, (simple_sel, ), more_selectors, _, block = children
    tag = tag[0] if tag else None

    selectors = [simple_sel]
    if more_selectors:
        selectors = selectors.extend([selector for _, (selector, ) in more_selectors])

    return RuleBlock(tag=tag,
                     selectors=selectors,
                     block=block,
                     index=node.start)

### AT-RULES ###

@GDef(r'at_rule = tag? "@" (at_rule_keyframe / at_rule_normal)')
def at_rule(env, node, children):
    tag = children[0][0] if children[0] else None
    text, block = children[2][0]

    return AtRuleBlock(tag=tag,
                       text=text.strip(),
                       block=block,
                       index=node.start)

@GDef(r'at_rule_normal = any ___ block')
def at_rule_normal(env, node, children):
    text, _, block = children
    return (text, block)

@GDef(r'at_rule_keyframe = "keyframes" _ name ___ keyframes_block')
def at_rule_keyframes(env, node, children):
    _, _, name, ___, block = children
    return ("keyframes " + name, block)

@GDef(r'keyframes_block = blk_open ___ keyframe+ blk_close')
def keyframes_block(env, node, children):
    return children[2]

@GDef(r'keyframe = tag? expr ___ block ___')
def keyframe(env, node, children):
    tag, condition, _, block, _ = children
    tag = tag[0] if tag else None
    return Keyframe(
        tag=tag,
        condition=condition,
        block=block,
        index=node.start)

### CSS RULES ###

@GDef(r'declaration = tag? property ":" _ expr+')
def declaration(env, node, children):
    tag, prop, _, _, values = children
    tag = tag[0] if tag else None
    return CSSRule(tag=tag,
                   prop=prop,
                   values=values,
                   index=node.start)

@GDef(r'property = ~"[a-zA-Z\_][a-zA-Z0-9\-\_]*"')
def property_(env, node, children):
    return node.text

@GDef(r'expr = (mixin_or_func_call / math / color / string_val) _?')
def expr(env, node, children):
    return children[0][0]

### MIXINS, FUNCTIONS, and VARIABLES ###

@GDef(r'mixin_decl = name "(" (_? name _?)* "):" ___ block',
      defer=True)
def mixin_decl(env, node):
    name, _, variables, _, _, block = node

    ly_engine = LY(env=env)
    name = ly_engine._evalnode(name)
    variables = [varname for _, varname, _ in ly_engine._evalnode(variables)]

    def f(*args):
        global_vars = list(env.items())
        local_vars = list(zip(variables, args))
        temp_env = dict(global_vars + local_vars)

        ly_engine_local = LY(env=temp_env)
        return ly_engine_local._evalnode(block)

    env[name] = MixIn(name=name, func=f)

@GDef(r'mixin_or_func_call = name "(" (_? expr _?)* ")"')
def mixin_or_func_call(env, node, children):
    name, _, args, _ = children
    args = [arg for _, arg, _ in args]

    if name in builtin_funcs:
        return builtin_funcs[name](*args)
    elif name in env:
        return UnpackMe(env[name](*args))

    raise UnknownMixinOrFunc(location=node.start)

@GDef(r'lvalue = ~"[a-zA-Z\_][a-zA-Z0-9\-\_]*"')
def lvalue(env, node, children):
    if node.text in env:
        return env[node.text]
    return node.text

@GDef(r'var_decl = name _? "=" _? expr ___')
def var_decl(env, node, children):
    name, _, _, _, value, _ = children
    env[name] = value

### ROOT BLOCKS ###

@GDef(r'root_block = "(root" ("." name)? ")" _? string_val?')
def root_block(env, node, children):
    _, possible_name, _, _, prefix = children
    tag_name = possible_name[0][1] if possible_name else None
    if not prefix:
        prefix = ''
    else:
        prefix = prefix[0]
    return RootBlock(tag_name=tag_name, prefix=prefix)

### TEXT ###

@GDef(r'''string_val = ~'\"[^"\n]*\"' ''')
def string_val(env, node, children):
    return node.text[1:-1]

### NUMERIC ###

@GDef(r'math = sum')
def math(env, node, children):
    return children[0]

@GDef(r'''
sum = prod (_? sum_op _? prod)*
prod = equality (_ prod_op _ equality)*
equality = value (_ equality_op _ value)*
''')
def math_operation(env, node, children):
    if len(children[1]):
        operations = [
            (operator, operand) for _, operator, _, operand in children[1]]
        return do_math(children[0], operations)
    return children[0]

@GDef(r'''
sum_op = "+" / "-"
prod_op = "*" / "/"
equality_op = "=="
''')
def operator_symbols(env, node, children):
    return node.text

@GDef(r'value = num_val / lvalue / paren')
def math_value(env, node, children):
    return children[0]

@GDef(r'paren = "(" _? sum _? ")"')
def math_paren(env, node, children):
    return children[2]

@GDef(r'num_val = distance / num / hexcolor / color_name / lvalue')
def num_val(env, node, children):
    return children[0]

@GDef(r'distance = num unit')
def distance(env, node, children):
    num, unit = children
    return Distance(value=num, unit=unit)

@GDef(r'num = ~"\-?\d+(\.\d+)?"')
def num(env, node, children):
    return Decimal(node.text)

@GDef(r'unit = "em" / "ex" / "px" / "cm" / "mm" / "in" / "pt" / "pc" / "deg" / "rad" / "grad" / "ms" / "s" / "hz" / "khz" / "%"')
def unit(env, node, children):
    return node.text

### MATH HELPER FUNCTION ###

OPERATORS = {'+': op.add,
             '-': op.sub,
             '*': op.mul,
             '/': op.truediv,
             '==': op.eq
             }

def do_math(start, operations):
    output_val = start
    for optr_symbol, operand in operations:
        output_val = OPERATORS[optr_symbol](output_val, operand)
    return output_val

### COLORS ###

@GDef(r'color = color_name / hexcolor')
def color(env, node, children):
    return node.text

@GDef(r'hexcolor = "#" ~"[0-9a-fA-F]+"')
def color_(env, node, children):
    return Color(node.text, ctype='hex')

@GDef(r'color_name = "aliceblue" / "antiquewhite" / "aqua" / "aquamarine" / "azure" / "beige" / "bisque" / "black" / "blanchedalmond" / "blue" / "blueviolet" / "brown" / "burlywood" / "cadetblue" / "chartreuse" / "chocolate" / "coral" / "cornflowerblue" / "cornsilk" / "crimson" / "cyan" / "darkblue" / "darkcyan" / "darkgoldenrod" / "darkgray" / "darkgreen" / "darkkhaki" / "darkmagenta" / "darkolivegreen" / "darkorange" / "darkorchid" / "darkred" / "darksalmon" / "darkseagreen" / "darkslateblue" / "darkslategray" / "darkturquoise" / "darkviolet" / "deeppink" / "deepskyblue" / "dimgray" / "dodgerblue" / "firebrick" / "floralwhite" / "forestgreen" / "fuchsia" / "gainsboro" / "ghostwhite" / "gold" / "goldenrod" / "gray" / "green" / "greenyellow" / "honeydew" / "hotpink" / "indianred " / "indigo " / "ivory" / "khaki" / "lavender" / "lavenderblush" / "lawngreen" / "lemonchiffon" / "lightblue" / "lightcoral" / "lightcyan" / "lightgoldenrodyellow" / "lightgray" / "lightgreen" / "lightpink" / "lightsalmon" / "lightseagreen" / "lightskyblue" / "lightslategray" / "lightsteelblue" / "lightyellow" / "lime" / "limegreen" / "linen" / "magenta" / "maroon" / "mediumaquamarine" / "mediumblue" / "mediumorchid" / "mediumpurple" / "mediumseagreen" / "mediumslateblue" / "mediumspringgreen" / "mediumturquoise" / "mediumvioletred" / "midnightblue" / "mintcream" / "mistyrose" / "moccasin" / "navajowhite" / "navy" / "oldlace" / "olive" / "olivedrab" / "orange" / "orangered" / "orchid" / "palegoldenrod" / "palegreen" / "paleturquoise" / "palevioletred" / "papayawhip" / "peachpuff" / "peru" / "pink" / "plum" / "powderblue" / "purple" / "red" / "rosybrown" / "royalblue" / "saddlebrown" / "salmon" / "sandybrown" / "seagreen" / "seashell" / "sienna" / "silver" / "skyblue" / "slateblue" / "slategray" / "snow" / "springgreen" / "steelblue" / "tan" / "teal" / "thistle" / "tomato" / "turquoise" / "violet" / "wheat" / "white" / "whitesmoke" / "yellow" / "yellowgreen"')
def color_name(env, node, children):
    return Color(node.text, ctype='named')

### BASIC ###

@GDef(r'tag = "(" name ")" _')
def tag(env, node, children):
    return children[1]

@GDef(r'name = ~"[a-zA-Z\_][a-zA-Z0-9\-\_]*"')
def name(env, node, children):
    return node.text

@GDef(r'any = ~".+"')
def any(env, node, children):
    return node.text.strip()

@GDef(r'''
___ = ~"[\s]*" (comment ~"[\s]*")*
_ = ~"[ \t]+"
blk_open = "{blk_open}"
blk_close = "{blk_close}"
comment = ~"{comment_open}.*(?!{comment_close})"
'''.format(
    comment_open=COMMENT_OPEN,
    comment_close=COMMENT_CLOSE,
    blk_open=BLK_OPEN,
    blk_close=BLK_CLOSE)
)
def return_none(env, node, children):
    return None

from .core import LY
<|endoftext|>


<commit_before>"""
Module to call mcsqs, distributed with AT-AT
https://www.brown.edu/Departments/Engineering/Labs/avdw/atat/
"""

import os
import warnings
from subprocess import Popen, TimeoutExpired
from typing import Dict, Union, List, NamedTuple, Optional
from typing_extensions import Literal
from pathlib import Path

from monty.dev import requires
from monty.os.path import which
import tempfile

from pymatgen import Structure

class Sqs(NamedTuple):
    """
    Return type for run_mcsqs.
    """

    bestsqs: Structure
    objective_function: Union[float, Literal["Perfect_match"]]
    allsqs: List
    clusters: List
    directory: str

@requires(
    which("mcsqs") and which("str2cif"),
    "run_mcsqs requires first installing AT-AT, "
    "see https://www.brown.edu/Departments/Engineering/Labs/avdw/atat/",
)
def run_mcsqs(
    structure: Structure,
    clusters: Dict[int, float],
    scaling: Union[int, List[int]],
    search_time: float = 60,
    directory: Optional[str] = None,
    instances: Optional[int] = None,
    temperature: Union[int, float] = 1,
    wr: float = 1,
    wn: float = 1,
    wd: float = 0.5,
    tol: float = 1e-3,
) -> Sqs:
    """
    Helper function for calling mcsqs with different arguments
    Args:
        structure (Structure): Disordered pymatgen Structure object
        clusters (dict): Dictionary of cluster interactions with entries in the form
            number of atoms: cutoff in angstroms
        scaling (int or list): Scaling factor to determine supercell. Two options are possible:
                a. (preferred) Scales number of atoms, e.g., for a structure with 8 atoms,
                   scaling=4 would lead to a 32 atom supercell
                b. A sequence of three scaling factors, e.g., [2, 1, 1], which
                   specifies that the supercell should have dimensions 2a x b x c
    Keyword Args:
        search_time (float): Time spent looking for the ideal SQS in minutes (default: 60)
        directory (str): Directory to run mcsqs calculation and store files (default: None
            runs calculations in a temp directory)
        instances (int): Specifies the number of parallel instances of mcsqs to run
            (default: number of cpu cores detected by Python)
        temperature (int or float): Monte Carlo temperature (default: 1), "T" in atat code
        wr (int or float): Weight assigned to range of perfect correlation match in objective
            function (default = 1)
        wn (int or float): Multiplicative decrease in weight per additional point in cluster (default: 1)
        wd (int or float): Exponent of decay in weight as function of cluster diameter (default: 0.5)
        tol (int or float): Tolerance for matching correlations (default: 1e-3)

    Returns:
        Tuple of Pymatgen structure SQS of the input structure, the mcsqs objective function,
            list of all SQS structures, and the directory where calculations are run
    """

    num_atoms = len(structure)

    if structure.is_ordered:
        raise ValueError("Pick a disordered structure")

    if instances is None:
        # os.cpu_count() can return None if detection fails
        instances = os.cpu_count()

    original_directory = os.getcwd()
    if not directory:
        directory = tempfile.mkdtemp()
    os.chdir(directory)

    if isinstance(scaling, (int, float)):

        if scaling % 1:
            raise ValueError("Scaling should be an integer, not {}".format(scaling))
        mcsqs_find_sqs_cmd = ["mcsqs", "-n {}".format(scaling * num_atoms)]

    else:

        # Set supercell to identity (will make supercell with pymatgen)
        with open("sqscell.out", "w") as f:
            f.write("1\n1 0 0\n0 1 0\n0 0 1\n")
        structure = structure * scaling
        mcsqs_find_sqs_cmd = ["mcsqs", "-rc", "-n {}".format(num_atoms)]

    structure.to(filename="rndstr.in")

    # Generate clusters
    mcsqs_generate_clusters_cmd = ["mcsqs"]
    for num in clusters:
        mcsqs_generate_clusters_cmd.append("-" + str(num) + "=" + str(clusters[num]))

    # Run mcsqs to find clusters
    p = Popen(mcsqs_generate_clusters_cmd)
    p.communicate()

    # Generate SQS structures
    add_ons = [
        "-T {}".format(temperature),
        "-wr {}".format(wr),
        "-wn {}".format(wn),
        "-wd {}".format(wd),
        "-tol {}".format(tol),
    ]

    mcsqs_find_sqs_processes = []
    if instances and instances > 1:
        # if multiple instances, run a range of commands using "-ip"
        for i in range(instances):
            instance_cmd = ["-ip {}".format(i + 1)]
            cmd = mcsqs_find_sqs_cmd + add_ons + instance_cmd
            p = Popen(cmd)
            mcsqs_find_sqs_processes.append(p)
    else:
        # run normal mcsqs command
        cmd = mcsqs_find_sqs_cmd + add_ons
        p = Popen(cmd)
        mcsqs_find_sqs_processes.append(p)

    try:
        for idx, p in enumerate(mcsqs_find_sqs_processes):
            p.communicate(timeout=search_time * 60)

        if instances and instances > 1:
            p = Popen(["mcsqs", "-best"])
            p.communicate()

        if os.path.exists("bestsqs.out") and os.path.exists("bestcorr.out"):
            return _parse_sqs_path(".")

        raise RuntimeError("mcsqs exited before timeout reached")

    except TimeoutExpired:
        for p in mcsqs_find_sqs_processes:
            p.kill()
            p.communicate()

        # Find the best sqs structures
        if instances and instances > 1:

            if not os.path.exists("bestcorr1.out"):
                raise RuntimeError(
                    "mcsqs did not generate output files, "
                    "is search_time sufficient or are number of instances too high?"
                )

            p = Popen(["mcsqs", "-best"])
            p.communicate()

        if os.path.exists("bestsqs.out") and os.path.exists("bestcorr.out"):
            sqs = _parse_sqs_path(".")
            return sqs

        else:
            os.chdir(original_directory)
            raise TimeoutError("Cluster expansion took too long.")

def _parse_sqs_path(path) -> Sqs:
    """
    Private function to parse mcsqs output directory
    Args:
        path: directory to perform parsing

    Returns:
        Tuple of Pymatgen structure SQS of the input structure, the mcsqs objective function,
            list of all SQS structures, and the directory where calculations are run
    """

    path = Path(path)

    # detected instances will be 0 if mcsqs was run in series, or number of instances
    detected_instances = len(list(path.glob("bestsqs*[0-9]*.out")))

    # Convert best SQS structure to cif file and pymatgen Structure
    p = Popen("str2cif < bestsqs.out > bestsqs.cif", shell=True, cwd=path)
    p.communicate()

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        bestsqs = Structure.from_file(path / "bestsqs.cif")

    # Get best SQS objective function
    with open(path / "bestcorr.out", "r") as f:
        lines = f.readlines()

    objective_function_str = lines[-1].split("=")[-1].strip()
    objective_function: Union[float, Literal["Perfect_match"]]
    if objective_function_str != "Perfect_match":
        objective_function = float(objective_function_str)
    else:
        objective_function = "Perfect_match"

    # Get all SQS structures and objective functions
    allsqs = []

    for i in range(detected_instances):
        sqs_out = "bestsqs{}.out".format(i + 1)
        sqs_cif = "bestsqs{}.cif".format(i + 1)
        corr_out = "bestcorr{}.out".format(i + 1)
        p = Popen("str2cif < {} > {}".format(sqs_out, sqs_cif), shell=True, cwd=path)
        p.communicate()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            sqs = Structure.from_file(path / sqs_cif)
        with open(path / corr_out, "r") as f:
            lines = f.readlines()

        objective_function_str = lines[-1].split("=")[-1].strip()
        obj: Union[float, Literal["Perfect_match"]]
        if objective_function_str != "Perfect_match":
            obj = float(objective_function_str)
        else:
            obj = "Perfect_match"
        allsqs.append({"structure": sqs, "objective_function": obj})

    clusters = _parse_clusters(path / "clusters.out")

    return Sqs(
        bestsqs=bestsqs,
        objective_function=objective_function,
        allsqs=allsqs,
        directory=str(path.resolve()),
        clusters=clusters
    )

def _parse_clusters(filename):
    """
    Private function to parse clusters.out file
    Args:
        path: directory to perform parsing

    Returns:
        List of dicts
    """

    with open(filename, "r") as f:
        lines = f.readlines()

    clusters = []
    cluster_block = []
    for line in lines:
        line = line.split("\n")[0]
        if line == '':
            clusters.append(cluster_block)
            cluster_block = []
        else:
            cluster_block.append(line)

    cluster_dicts = []
    for cluster in clusters:
        cluster_dict = {"multiplicity": int(cluster[0]),
                        "longest_pair_length": float(cluster[1]),
                        "num_points_in_cluster": int(cluster[2])}
        points = []
        for point in range(cluster_dict["num_points_in_cluster"]):
            line = cluster[3 + point].split(" ")
            point_dict = {}
            point_dict["coordinates"] = [float(line) for line in line[0:3]]
            point_dict["num_possible_species_minus_2"] = int(line[3])
            point_dict["cluster_function"] = float(line[4])
            points.append(point_dict)

        cluster_dict["coordinates"] = points
        cluster_dicts.append(cluster_dict)

    return cluster_dicts
<commit_msg>Parse structure from .out file<commit_after>"""
Module to call mcsqs, distributed with AT-AT
https://www.brown.edu/Departments/Engineering/Labs/avdw/atat/
"""

import os
import warnings
from subprocess import Popen, TimeoutExpired
from typing import Dict, Union, List, NamedTuple, Optional
from typing_extensions import Literal
from pathlib import Path

from monty.dev import requires
from monty.os.path import which
import tempfile

from pymatgen import Structure

class Sqs(NamedTuple):
    """
    Return type for run_mcsqs.
    """

    bestsqs: Structure
    objective_function: Union[float, Literal["Perfect_match"]]
    allsqs: List
    clusters: List
    directory: str

@requires(
    which("mcsqs") and which("str2cif"),
    "run_mcsqs requires first installing AT-AT, "
    "see https://www.brown.edu/Departments/Engineering/Labs/avdw/atat/",
)
def run_mcsqs(
    structure: Structure,
    clusters: Dict[int, float],
    scaling: Union[int, List[int]],
    search_time: float = 60,
    directory: Optional[str] = None,
    instances: Optional[int] = None,
    temperature: Union[int, float] = 1,
    wr: float = 1,
    wn: float = 1,
    wd: float = 0.5,
    tol: float = 1e-3,
) -> Sqs:
    """
    Helper function for calling mcsqs with different arguments
    Args:
        structure (Structure): Disordered pymatgen Structure object
        clusters (dict): Dictionary of cluster interactions with entries in the form
            number of atoms: cutoff in angstroms
        scaling (int or list): Scaling factor to determine supercell. Two options are possible:
                a. (preferred) Scales number of atoms, e.g., for a structure with 8 atoms,
                   scaling=4 would lead to a 32 atom supercell
                b. A sequence of three scaling factors, e.g., [2, 1, 1], which
                   specifies that the supercell should have dimensions 2a x b x c
    Keyword Args:
        search_time (float): Time spent looking for the ideal SQS in minutes (default: 60)
        directory (str): Directory to run mcsqs calculation and store files (default: None
            runs calculations in a temp directory)
        instances (int): Specifies the number of parallel instances of mcsqs to run
            (default: number of cpu cores detected by Python)
        temperature (int or float): Monte Carlo temperature (default: 1), "T" in atat code
        wr (int or float): Weight assigned to range of perfect correlation match in objective
            function (default = 1)
        wn (int or float): Multiplicative decrease in weight per additional point in cluster (default: 1)
        wd (int or float): Exponent of decay in weight as function of cluster diameter (default: 0.5)
        tol (int or float): Tolerance for matching correlations (default: 1e-3)

    Returns:
        Tuple of Pymatgen structure SQS of the input structure, the mcsqs objective function,
            list of all SQS structures, and the directory where calculations are run
    """

    num_atoms = len(structure)

    if structure.is_ordered:
        raise ValueError("Pick a disordered structure")

    if instances is None:
        # os.cpu_count() can return None if detection fails
        instances = os.cpu_count()

    original_directory = os.getcwd()
    if not directory:
        directory = tempfile.mkdtemp()
    os.chdir(directory)

    if isinstance(scaling, (int, float)):

        if scaling % 1:
            raise ValueError("Scaling should be an integer, not {}".format(scaling))
        mcsqs_find_sqs_cmd = ["mcsqs", "-n {}".format(scaling * num_atoms)]

    else:

        # Set supercell to identity (will make supercell with pymatgen)
        with open("sqscell.out", "w") as f:
            f.write("1\n1 0 0\n0 1 0\n0 0 1\n")
        structure = structure * scaling
        mcsqs_find_sqs_cmd = ["mcsqs", "-rc", "-n {}".format(num_atoms)]

    structure.to(filename="rndstr.in")

    # Generate clusters
    mcsqs_generate_clusters_cmd = ["mcsqs"]
    for num in clusters:
        mcsqs_generate_clusters_cmd.append("-" + str(num) + "=" + str(clusters[num]))

    # Run mcsqs to find clusters
    p = Popen(mcsqs_generate_clusters_cmd)
    p.communicate()

    # Generate SQS structures
    add_ons = [
        "-T {}".format(temperature),
        "-wr {}".format(wr),
        "-wn {}".format(wn),
        "-wd {}".format(wd),
        "-tol {}".format(tol),
    ]

    mcsqs_find_sqs_processes = []
    if instances and instances > 1:
        # if multiple instances, run a range of commands using "-ip"
        for i in range(instances):
            instance_cmd = ["-ip {}".format(i + 1)]
            cmd = mcsqs_find_sqs_cmd + add_ons + instance_cmd
            p = Popen(cmd)
            mcsqs_find_sqs_processes.append(p)
    else:
        # run normal mcsqs command
        cmd = mcsqs_find_sqs_cmd + add_ons
        p = Popen(cmd)
        mcsqs_find_sqs_processes.append(p)

    try:
        for idx, p in enumerate(mcsqs_find_sqs_processes):
            p.communicate(timeout=search_time * 60)

        if instances and instances > 1:
            p = Popen(["mcsqs", "-best"])
            p.communicate()

        if os.path.exists("bestsqs.out") and os.path.exists("bestcorr.out"):
            return _parse_sqs_path(".")

        raise RuntimeError("mcsqs exited before timeout reached")

    except TimeoutExpired:
        for p in mcsqs_find_sqs_processes:
            p.kill()
            p.communicate()

        # Find the best sqs structures
        if instances and instances > 1:

            if not os.path.exists("bestcorr1.out"):
                raise RuntimeError(
                    "mcsqs did not generate output files, "
                    "is search_time sufficient or are number of instances too high?"
                )

            p = Popen(["mcsqs", "-best"])
            p.communicate()

        if os.path.exists("bestsqs.out") and os.path.exists("bestcorr.out"):
            sqs = _parse_sqs_path(".")
            return sqs

        else:
            os.chdir(original_directory)
            raise TimeoutError("Cluster expansion took too long.")

def _parse_sqs_path(path) -> Sqs:
    """
    Private function to parse mcsqs output directory
    Args:
        path: directory to perform parsing

    Returns:
        Tuple of Pymatgen structure SQS of the input structure, the mcsqs objective function,
            list of all SQS structures, and the directory where calculations are run
    """

    path = Path(path)

    # detected instances will be 0 if mcsqs was run in series, or number of instances
    detected_instances = len(list(path.glob("bestsqs*[0-9]*.out")))

    # Convert best SQS structure to cif file and pymatgen Structure
    p = Popen("str2cif < bestsqs.out > bestsqs.cif", shell=True, cwd=path)
    p.communicate()

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        bestsqs = Structure.from_file(path / "bestsqs.out")

    # Get best SQS objective function
    with open(path / "bestcorr.out", "r") as f:
        lines = f.readlines()

    objective_function_str = lines[-1].split("=")[-1].strip()
    objective_function: Union[float, Literal["Perfect_match"]]
    if objective_function_str != "Perfect_match":
        objective_function = float(objective_function_str)
    else:
        objective_function = "Perfect_match"

    # Get all SQS structures and objective functions
    allsqs = []

    for i in range(detected_instances):
        sqs_out = "bestsqs{}.out".format(i + 1)
        sqs_cif = "bestsqs{}.cif".format(i + 1)
        corr_out = "bestcorr{}.out".format(i + 1)
        p = Popen("str2cif < {} > {}".format(sqs_out, sqs_cif), shell=True, cwd=path)
        p.communicate()
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            sqs = Structure.from_file(path / sqs_out)
        with open(path / corr_out, "r") as f:
            lines = f.readlines()

        objective_function_str = lines[-1].split("=")[-1].strip()
        obj: Union[float, Literal["Perfect_match"]]
        if objective_function_str != "Perfect_match":
            obj = float(objective_function_str)
        else:
            obj = "Perfect_match"
        allsqs.append({"structure": sqs, "objective_function": obj})

    clusters = _parse_clusters(path / "clusters.out")

    return Sqs(
        bestsqs=bestsqs,
        objective_function=objective_function,
        allsqs=allsqs,
        directory=str(path.resolve()),
        clusters=clusters
    )

def _parse_clusters(filename):
    """
    Private function to parse clusters.out file
    Args:
        path: directory to perform parsing

    Returns:
        List of dicts
    """

    with open(filename, "r") as f:
        lines = f.readlines()

    clusters = []
    cluster_block = []
    for line in lines:
        line = line.split("\n")[0]
        if line == '':
            clusters.append(cluster_block)
            cluster_block = []
        else:
            cluster_block.append(line)

    cluster_dicts = []
    for cluster in clusters:
        cluster_dict = {"multiplicity": int(cluster[0]),
                        "longest_pair_length": float(cluster[1]),
                        "num_points_in_cluster": int(cluster[2])}
        points = []
        for point in range(cluster_dict["num_points_in_cluster"]):
            line = cluster[3 + point].split(" ")
            point_dict = {}
            point_dict["coordinates"] = [float(line) for line in line[0:3]]
            point_dict["num_possible_species_minus_2"] = int(line[3])
            point_dict["cluster_function"] = float(line[4])
            points.append(point_dict)

        cluster_dict["coordinates"] = points
        cluster_dicts.append(cluster_dict)

    return cluster_dicts
<|endoftext|>


<commit_before>import datetime
import hashlib
import os
import time
import traceback

from .DisassemblyStatistics import DisassemblyStatistics
from .intel.IntelDisassembler import IntelDisassembler
from .ida.IdaExporter import IdaExporter
from smda.utility.FileLoader import FileLoader

class Disassembler(object):

    def __init__(self, config, backend="intel"):
        self.config = config
        self.disassembler = None
        if backend == "intel":
            self.disassembler = IntelDisassembler(config)
        elif backend == "IDA":
            self.disassembler = IdaExporter(config)
        self.disassembly = None
        self._start_time = None
        self._timeout = 0

    def _callbackAnalysisTimeout(self):
        if not self._timeout:
            return False
        return time.time() - self._start_time > self._timeout

    def disassembleFile(self, file_path, pdb_path=""):
        loader = FileLoader(file_path, map_file=True)
        base_addr = loader.getBaseAddress()
        file_content = loader.getData()
        start = time.clock()
        try:
            self.disassembler.setFilePath(file_path)
            self.disassembler.addPdbFile(pdb_path, base_addr)
            disassembly = self.disassemble(file_content, base_addr, timeout=self.config.TIMEOUT)
            report = self.getDisassemblyReport(disassembly)
            report["filename"] = os.path.basename(file_path)
            print(disassembly)
        except Exception as exc:
            print("-> an error occured (", str(exc), ").")
            report = {"status":"error", "meta": {"traceback": traceback.format_exc(exc)}, "execution_time": time.clock() - start}
        return report

    def disassembleBuffer(self, file_content, base_addr, bitness=None):
        start = time.clock()
        try:
            self.disassembler.setFilePath("")
            disassembly = self.disassemble(file_content, base_addr, bitness, timeout=self.config.TIMEOUT)
            report = self.getDisassemblyReport(disassembly)
            report["filename"] = ""
            print(disassembly)
        except Exception as exc:
            print("-> an error occured (", str(exc), ").")
            report = {"status":"error", "meta": {"traceback": traceback.format_exc(exc)}, "execution_time": time.clock() - start}
        return report

    def disassemble(self, binary, base_addr, bitness=None, timeout=0):
        self._start_time = time.time()
        self._timeout = timeout
        self.disassembly = self.disassembler.analyzeBuffer(binary, base_addr, bitness, self._callbackAnalysisTimeout)
        return self.disassembly

    def getDisassemblyReport(self, disassembly=None):
        report = {}
        if disassembly is None:
            if self.disassembly is not None:
                disassembly = self.disassembly
            else:
                return {}
        stats = DisassemblyStatistics(disassembly)
        report = {
            "architecture": disassembly.architecture,
            "base_addr": disassembly.base_addr,
            "bitness": disassembly.bitness,
            "buffer_size": len(disassembly.binary),
            "disassembly_errors": disassembly.errors,
            "execution_time": disassembly.getAnalysisDuration(),
            "metadata" : {
                "message": "Analysis finished regularly."
            },
            "sha256": hashlib.sha256(disassembly.binary).hexdigest(),
            "smda_version": self.config.VERSION,
            "status": disassembly.getAnalysisOutcome(),
            "summary": stats.calculate(),
            "timestamp": datetime.datetime.utcnow().strftime("%Y-%m-%dT%H-%M-%S"),
            "xcfg": disassembly.collectCfg(),
        }
        return report
<commit_msg>removed deprecated time.clock()<commit_after>import datetime
import hashlib
import os
import traceback

from .DisassemblyStatistics import DisassemblyStatistics
from .intel.IntelDisassembler import IntelDisassembler
from .ida.IdaExporter import IdaExporter
from smda.utility.FileLoader import FileLoader

class Disassembler(object):

    def __init__(self, config, backend="intel"):
        self.config = config
        self.disassembler = None
        if backend == "intel":
            self.disassembler = IntelDisassembler(config)
        elif backend == "IDA":
            self.disassembler = IdaExporter(config)
        self.disassembly = None
        self._start_time = None
        self._timeout = 0

    def _getDurationInSeconds(self, start_ts, end_ts):
        return (self.analysis_end_ts - self.analysis_start_ts).seconds + ((self.analysis_end_ts - self.analysis_start_ts).microseconds / 1000000.0)

    def _callbackAnalysisTimeout(self):
        if not self._timeout:
            return False
        time_diff = datetime.datetime.utcnow() - self._start_time
        return time_diff.seconds >= self._timeout

    def disassembleFile(self, file_path, pdb_path=""):
        loader = FileLoader(file_path, map_file=True)
        base_addr = loader.getBaseAddress()
        file_content = loader.getData()
        start = datetime.datetime.utcnow()
        try:
            self.disassembler.setFilePath(file_path)
            self.disassembler.addPdbFile(pdb_path, base_addr)
            disassembly = self.disassemble(file_content, base_addr, timeout=self.config.TIMEOUT)
            report = self.getDisassemblyReport(disassembly)
            report["filename"] = os.path.basename(file_path)
            print(disassembly)
        except Exception as exc:
            print("-> an error occured (", str(exc), ").")
            report = {"status":"error", "meta": {"traceback": traceback.format_exc(exc)}, "execution_time": self._getDurationInSeconds(start, datetime.datetime.utcnow())}
        return report

    def disassembleBuffer(self, file_content, base_addr, bitness=None):
        start = datetime.datetime.utcnow()
        try:
            self.disassembler.setFilePath("")
            disassembly = self.disassemble(file_content, base_addr, bitness, timeout=self.config.TIMEOUT)
            report = self.getDisassemblyReport(disassembly)
            report["filename"] = ""
            print(disassembly)
        except Exception as exc:
            print("-> an error occured (", str(exc), ").")
            report = {"status":"error", "meta": {"traceback": traceback.format_exc(exc)}, "execution_time": self._getDurationInSeconds(start, datetime.datetime.utcnow())}
        return report

    def disassemble(self, binary, base_addr, bitness=None, timeout=0):
        self._start_time = datetime.datetime.utcnow()
        self._timeout = timeout
        self.disassembly = self.disassembler.analyzeBuffer(binary, base_addr, bitness, self._callbackAnalysisTimeout)
        return self.disassembly

    def getDisassemblyReport(self, disassembly=None):
        report = {}
        if disassembly is None:
            if self.disassembly is not None:
                disassembly = self.disassembly
            else:
                return {}
        stats = DisassemblyStatistics(disassembly)
        report = {
            "architecture": disassembly.architecture,
            "base_addr": disassembly.base_addr,
            "bitness": disassembly.bitness,
            "buffer_size": len(disassembly.binary),
            "disassembly_errors": disassembly.errors,
            "execution_time": disassembly.getAnalysisDuration(),
            "metadata" : {
                "message": "Analysis finished regularly."
            },
            "sha256": hashlib.sha256(disassembly.binary).hexdigest(),
            "smda_version": self.config.VERSION,
            "status": disassembly.getAnalysisOutcome(),
            "summary": stats.calculate(),
            "timestamp": datetime.datetime.utcnow().strftime("%Y-%m-%dT%H-%M-%S"),
            "xcfg": disassembly.collectCfg(),
        }
        return report
<|endoftext|>


<commit_before>
import numpy as np
from scipy import stats
import pandas as pd

def beta_geometric_nbd_model(T, r, alpha, a, b, size=1):
    """
    Generate artificial data according to the BG/NBD model. See [1] for model details

    Parameters:
        T: scalar or array, the length of time observing new customers.
        r, alpha, a, b: scalars, represening parameters in the model. See [1]
        size: the number of customers to generate

    Returns:
        DataFrame, with index as customer_ids and the following columns:
        'frequency', 'recency', 'T', 'lambda', 'p', 'alive', 'customer_id'

    [1]: '"Counting Your Customers" the Easy Way: An Alternative to the Pareto/NBD Model'
    (http://brucehardie.com/papers/bgnbd_2004-04-20.pdf)

    """

    probability_of_post_purchase_death = stats.beta.rvs(a, b, size=size)
    lambda_ = stats.gamma.rvs(r, scale=1. / alpha, size=size)
    if type(T) in [float, int]:
        T = T * np.ones(size)
    else:
        T = np.asarray(T)

    columns = ['frequency', 'recency', 'T', 'lambda', 'p', 'alive', 'customer_id']
    df = pd.DataFrame(np.zeros((size, len(columns))), columns=columns)

    for i in xrange(size):
        p = probability_of_post_purchase_death[i]
        l = lambda_[i]
        churn = stats.geom.rvs(p)
        times = np.cumsum(np.r_[[0], stats.expon.rvs(scale=1. / l, size=churn)])
        t_cal = times[times < T[i]]
        df.ix[i] = len(t_cal) - 1,  max(t_cal), T[i], l, p, churn > len(t_cal) - 1, i

    return df.set_index('customer_id')
<commit_msg>this should be equivilant<commit_after>
import numpy as np
from scipy import stats
import pandas as pd

def beta_geometric_nbd_model(T, r, alpha, a, b, size=1):
    """
    Generate artificial data according to the BG/NBD model. See [1] for model details

    Parameters:
        T: scalar or array, the length of time observing new customers.
        r, alpha, a, b: scalars, represening parameters in the model. See [1]
        size: the number of customers to generate

    Returns:
        DataFrame, with index as customer_ids and the following columns:
        'frequency', 'recency', 'T', 'lambda', 'p', 'alive', 'customer_id'

    [1]: '"Counting Your Customers" the Easy Way: An Alternative to the Pareto/NBD Model'
    (http://brucehardie.com/papers/bgnbd_2004-04-20.pdf)

    """
    if type(T) in [float, int]:
        T = T * np.ones(size)
    else:
        T = np.asarray(T)

    probability_of_post_purchase_death = stats.beta.rvs(a, b, size=size)
    lambda_ = stats.gamma.rvs(r, scale=1. / alpha, size=size)

    columns = ['frequency', 'recency', 'T', 'lambda', 'p', 'alive', 'customer_id']
    df = pd.DataFrame(np.zeros((size, len(columns))), columns=columns)

    for i in xrange(size):
        p = probability_of_post_purchase_death[i]
        l = lambda_[i]
        churn = stats.geom.rvs(p)
        times = np.cumsum(stats.expon.rvs(scale=1. / l, size=churn))
        t_cal = times[times < T[i]]
        df.ix[i] = len(t_cal),  np.max(t_cal if t_cal.shape[0] > 0 else 0), T[i], l, p, churn > len(t_cal), i

    return df.set_index('customer_id')
<|endoftext|>


<commit_before>import collections
import os

from gather.analyze import Collector
from gather.transaction import (
    DryRunner, FilesystemTransaction, RollbackError
)
import gather.log as log

def recurse_file_iterator(roots):
    for path in roots:
        if os.path.isfile(path):
            yield path
        elif os.path.isdir(path):
            for container, _dirnames, filenames in os.walk(path):
                for filename in filenames:
                    yield os.path.join(container, filename)

DEFAULT_DIR_TEMPLATE = "{path_prefix}[{first}-{last}]{suffix}"

IGNORE = 0
REPORT = 1
CANCEL = 2
ALLOW = 3
SKIP = 4

AMBIGUITY_CHOICES = (IGNORE, REPORT, CANCEL)
SHARED_DIRECTORY_CHOICES = (ALLOW, SKIP, CANCEL)

CANCEL_REASON_AMBIGUITIES = 1
CANCEL_REASON_SHARED = 2

ROLLBACK_SET = 1
ROLLBACK_ALL = 2

RESULT_OK = 0
RESULT_CANCEL = 2
RESULT_ERROR_FULL_ROLLBACK = 3
RESULT_ERROR_PARTIAL_ROLLBACK = 4
RESULT_ERROR_FAILED_ROLLBACK = 5

MSG_CANCEL_REASONS = (
    (CANCEL_REASON_AMBIGUITIES, "ambiguous sequences"),
    (CANCEL_REASON_SHARED,      "multiple sequences sharing a directory"),
)
MSG_CANCEL_REASONS_REPORT = "Stopping because {reasons}"

MSG_AMBIGUOUS_HEADER = "The following files are ambiguous sequence members:"
MSG_AMBIGUOUS_NEXT = "  {file} could be followed by {choices[0]} or {choices[1]}"
MSG_AMBIGUOUS_PREVIOUS = "  {choices[0]} or {choices[1]} could precede {file}"

MSG_SHARED_HEADER_ALLOWED = "Directory will contain multiple sequences:"
MSG_SHARED_HEADER_DISALLOWED = "Directory would contain multiple sequences:"
MSG_SHARED_COACH = "Use the --template option to create distinct directory names, or allow directories to contain multiple sequences with --share allow"

MSG_SHORT = "Skipping sequence {sequence_name} because its length ({sequence_length}) is below the minimum {min_length}"

MSG_DRY_RUN = "--dry-run specified. No changes will be made."

MSG_ERROR = "Error: {error!s}. Rolling back..."
MSG_ROLLBACK_OK = "Rollback complete"
MSG_ROLLBACK_FAIL = """\
Error while rolling back - incomplete commands follow.
# The first command listed failed with exception: {error!s}"""
MSG_ROLLBACK_FAIL_EPILOG = "# End incomplete commands"

MSG_ROLLBACK_COUNT = "{count} of {total} sequences failed and were rolled back."

def gather(
    paths,
    recurse = False,
    dir_template = DEFAULT_DIR_TEMPLATE,
    min_sequence_length = 1,
    ambiguities = REPORT,
    shared_directories = ALLOW,
    rollback = ROLLBACK_ALL,
    dry_run = False,
):
    logger = log.Logger()
    collector = Collector()

    path_iter = recurse_file_iterator(paths) if recurse else paths
    collector.collect_all(path_iter)

    plan, cancel_reasons = prepare(
        collector,
        sequence_name_generator(dir_template),
        logger,
        min_sequence_length,
        ambiguities,
        shared_directories,
    )

    if dry_run:
        logger.info(MSG_DRY_RUN)
        fst = DryRunner()
    else:
        if len(cancel_reasons) > 0:
            return RESULT_CANCEL
        fst = FilesystemTransaction()

    rollbacks = 0
    for parent, paths in plan:
        logger.info(parent)
        try:
            fst.mkdirp(parent)
            for path in paths:
                logger.info("  %s" % path)
                new_path = os.path.join(parent, os.path.basename(path))
                fst.move(path, new_path)
            if rollback == ROLLBACK_SET:
                fst.commit()
            logger.info("")

        except OSError as ose:
            try:
                logger.error(MSG_ERROR, error=ose)
                fst.rollback()
                logger.info(MSG_ROLLBACK_OK)
                if rollback == ROLLBACK_ALL:
                    return RESULT_ERROR_FULL_ROLLBACK
                else:
                    rollbacks += 1

            except RollbackError as re:
                logger.critical(MSG_ROLLBACK_FAIL, error=re)
                for action in re.actions:
                    logger.critical("  %s" % action)
                logger.critical(MSG_ROLLBACK_FAIL_EPILOG)
                raise re

    if rollbacks != 0:
        logger.warning(MSG_ROLLBACK_COUNT, count=rollbacks, total=len(plan))
        if rollbacks == len(plan):
            return RESULT_ERROR_FULL_ROLLBACK
        return RESULT_ERROR_PARTIAL_ROLLBACK
    return RESULT_OK

AMB_TO_LOG_LEVEL = {
    IGNORE: log.VERBOSE,
    REPORT: log.INFO,
    CANCEL: log.ERROR,
}
SHARE_TO_LOG_LEVEL = {
    ALLOW:  log.VERBOSE,
    SKIP:   log.WARNING,
    CANCEL: log.ERROR,
}
def prepare(
    collector,
    sequence_namer,
    logger,
    min_sequence_length,
    ambiguities,
    shared_directories,
):
    plan = [ ]
    cancel_reasons = set()

    amb_log = logger.level_func(AMB_TO_LOG_LEVEL[ambiguities])
    share_log = logger.level_func(SHARE_TO_LOG_LEVEL[shared_directories])

    if collector.has_ambiguities():
        if ambiguities == CANCEL:
            cancel_reasons.add(CANCEL_REASON_AMBIGUITIES)

        amb_log(MSG_AMBIGUOUS_HEADER)
        for amb in collector.ambiguities():
            amb_log(format_ambiguity(amb))
        amb_log("")

    sequences_by_parent = group(collector.sequences(), key=sequence_namer)
    found_shared = False
    for parent, sequences in sequences_by_parent:
        if len(sequences) > 1:
            share_log(
                MSG_SHARED_HEADER_ALLOWED
                if shared_directories == ALLOW
                else MSG_SHARED_HEADER_DISALLOWED
            )

            share_log("  %s" % parent)
            for sequence in sequences:
                share_log("    %s" % sequence)
            share_log("")

            if shared_directories != ALLOW:
                if not found_shared:
                    found_shared = True
                    if shared_directories == CANCEL:
                        cancel_reasons.add(CANCEL_REASON_SHARED)
                    logger.defer.info(MSG_SHARED_COACH)
                continue

        for sequence in sequences:
            if len(sequence.paths) < min_sequence_length:
                logger.info(
                    MSG_SHORT,
                    sequence_name = str(sequence),
                    sequence_length = len(sequence.paths),
                    min_length = min_sequence_length
                )
            else:
                plan.append((parent, sequence.paths))

    if len(cancel_reasons) > 0:
        logger.error(
            MSG_CANCEL_REASONS_REPORT,
            reasons = ", ".join(
                message for reason, message in MSG_CANCEL_REASONS
                if reason in cancel_reasons
            )
        )

    logger.defer.flush()

    return plan, cancel_reasons

def sequence_name_generator(template):
    def generate(sequence):
        return template.format(
            path_prefix = os.path.join(sequence.container, sequence.prefix),
            name_prefix = sequence.prefix,
            suffix = sequence.suffix,
            first = sequence.first.number,
            last = sequence.last.number,
            field = "#" * sequence.first.digit_count,
        )
    return generate

def format_ambiguity(amb):
    template = (
        MSG_AMBIGUOUS_PREVIOUS
        if amb.direction == Collector.AMBIGUITY_PREVIOUS
        else MSG_AMBIGUOUS_NEXT
    )
    return template.format(**amb._as_dict())

def set_only_contains(query_set, element):
    return len(query_set) == 1 and element in query_set

def group(iterable, key=None):
    if key is None:
        key = lambda m: m

    result = collections.OrderedDict()

    for m in iterable:
        k = key(m)
        if k not in result:
            result[k] = [ m ]
        else:
            result[k].append(m)

    for k, ms in result.items():
        yield k, ms
<commit_msg>remove unused<commit_after>import collections
import os

from gather.analyze import Collector
from gather.transaction import (
    DryRunner, FilesystemTransaction, RollbackError
)
import gather.log as log

def recurse_file_iterator(roots):
    for path in roots:
        if os.path.isfile(path):
            yield path
        elif os.path.isdir(path):
            for container, _dirnames, filenames in os.walk(path):
                for filename in filenames:
                    yield os.path.join(container, filename)

DEFAULT_DIR_TEMPLATE = "{path_prefix}[{first}-{last}]{suffix}"

IGNORE = 0
REPORT = 1
CANCEL = 2
ALLOW = 3
SKIP = 4

AMBIGUITY_CHOICES = (IGNORE, REPORT, CANCEL)
SHARED_DIRECTORY_CHOICES = (ALLOW, SKIP, CANCEL)

CANCEL_REASON_AMBIGUITIES = 1
CANCEL_REASON_SHARED = 2

ROLLBACK_SET = 1
ROLLBACK_ALL = 2

RESULT_OK = 0
RESULT_CANCEL = 2
RESULT_ERROR_FULL_ROLLBACK = 3
RESULT_ERROR_PARTIAL_ROLLBACK = 4
RESULT_ERROR_FAILED_ROLLBACK = 5

MSG_CANCEL_REASONS = (
    (CANCEL_REASON_AMBIGUITIES, "ambiguous sequences"),
    (CANCEL_REASON_SHARED,      "multiple sequences sharing a directory"),
)
MSG_CANCEL_REASONS_REPORT = "Stopping because {reasons}"

MSG_AMBIGUOUS_HEADER = "The following files are ambiguous sequence members:"
MSG_AMBIGUOUS_NEXT = "  {file} could be followed by {choices[0]} or {choices[1]}"
MSG_AMBIGUOUS_PREVIOUS = "  {choices[0]} or {choices[1]} could precede {file}"

MSG_SHARED_HEADER_ALLOWED = "Directory will contain multiple sequences:"
MSG_SHARED_HEADER_DISALLOWED = "Directory would contain multiple sequences:"
MSG_SHARED_COACH = "Use the --template option to create distinct directory names, or allow directories to contain multiple sequences with --share allow"

MSG_SHORT = "Skipping sequence {sequence_name} because its length ({sequence_length}) is below the minimum {min_length}"

MSG_DRY_RUN = "--dry-run specified. No changes will be made."

MSG_ERROR = "Error: {error!s}. Rolling back..."
MSG_ROLLBACK_OK = "Rollback complete"
MSG_ROLLBACK_FAIL = """\
Error while rolling back - incomplete commands follow.
# The first command listed failed with exception: {error!s}"""
MSG_ROLLBACK_FAIL_EPILOG = "# End incomplete commands"

MSG_ROLLBACK_COUNT = "{count} of {total} sequences failed and were rolled back."

def gather(
    paths,
    recurse = False,
    dir_template = DEFAULT_DIR_TEMPLATE,
    min_sequence_length = 1,
    ambiguities = REPORT,
    shared_directories = ALLOW,
    rollback = ROLLBACK_ALL,
    dry_run = False,
):
    logger = log.Logger()
    collector = Collector()

    path_iter = recurse_file_iterator(paths) if recurse else paths
    collector.collect_all(path_iter)

    plan, cancel_reasons = prepare(
        collector,
        sequence_name_generator(dir_template),
        logger,
        min_sequence_length,
        ambiguities,
        shared_directories,
    )

    if dry_run:
        logger.info(MSG_DRY_RUN)
        fst = DryRunner()
    else:
        if len(cancel_reasons) > 0:
            return RESULT_CANCEL
        fst = FilesystemTransaction()

    rollbacks = 0
    for parent, paths in plan:
        logger.info(parent)
        try:
            fst.mkdirp(parent)
            for path in paths:
                logger.info("  %s" % path)
                new_path = os.path.join(parent, os.path.basename(path))
                fst.move(path, new_path)
            if rollback == ROLLBACK_SET:
                fst.commit()
            logger.info("")

        except OSError as ose:
            try:
                logger.error(MSG_ERROR, error=ose)
                fst.rollback()
                logger.info(MSG_ROLLBACK_OK)
                if rollback == ROLLBACK_ALL:
                    return RESULT_ERROR_FULL_ROLLBACK
                else:
                    rollbacks += 1

            except RollbackError as re:
                logger.critical(MSG_ROLLBACK_FAIL, error=re)
                for action in re.actions:
                    logger.critical("  %s" % action)
                logger.critical(MSG_ROLLBACK_FAIL_EPILOG)
                raise re

    if rollbacks != 0:
        logger.warning(MSG_ROLLBACK_COUNT, count=rollbacks, total=len(plan))
        if rollbacks == len(plan):
            return RESULT_ERROR_FULL_ROLLBACK
        return RESULT_ERROR_PARTIAL_ROLLBACK
    return RESULT_OK

AMB_TO_LOG_LEVEL = {
    IGNORE: log.VERBOSE,
    REPORT: log.INFO,
    CANCEL: log.ERROR,
}
SHARE_TO_LOG_LEVEL = {
    ALLOW:  log.VERBOSE,
    SKIP:   log.WARNING,
    CANCEL: log.ERROR,
}
def prepare(
    collector,
    sequence_namer,
    logger,
    min_sequence_length,
    ambiguities,
    shared_directories,
):
    plan = [ ]
    cancel_reasons = set()

    amb_log = logger.level_func(AMB_TO_LOG_LEVEL[ambiguities])
    share_log = logger.level_func(SHARE_TO_LOG_LEVEL[shared_directories])

    if collector.has_ambiguities():
        if ambiguities == CANCEL:
            cancel_reasons.add(CANCEL_REASON_AMBIGUITIES)

        amb_log(MSG_AMBIGUOUS_HEADER)
        for amb in collector.ambiguities():
            amb_log(format_ambiguity(amb))
        amb_log("")

    sequences_by_parent = group(collector.sequences(), key=sequence_namer)
    found_shared = False
    for parent, sequences in sequences_by_parent:
        if len(sequences) > 1:
            share_log(
                MSG_SHARED_HEADER_ALLOWED
                if shared_directories == ALLOW
                else MSG_SHARED_HEADER_DISALLOWED
            )

            share_log("  %s" % parent)
            for sequence in sequences:
                share_log("    %s" % sequence)
            share_log("")

            if shared_directories != ALLOW:
                if not found_shared:
                    found_shared = True
                    if shared_directories == CANCEL:
                        cancel_reasons.add(CANCEL_REASON_SHARED)
                    logger.defer.info(MSG_SHARED_COACH)
                continue

        for sequence in sequences:
            if len(sequence.paths) < min_sequence_length:
                logger.info(
                    MSG_SHORT,
                    sequence_name = str(sequence),
                    sequence_length = len(sequence.paths),
                    min_length = min_sequence_length
                )
            else:
                plan.append((parent, sequence.paths))

    if len(cancel_reasons) > 0:
        logger.error(
            MSG_CANCEL_REASONS_REPORT,
            reasons = ", ".join(
                message for reason, message in MSG_CANCEL_REASONS
                if reason in cancel_reasons
            )
        )

    logger.defer.flush()

    return plan, cancel_reasons

def sequence_name_generator(template):
    def generate(sequence):
        return template.format(
            path_prefix = os.path.join(sequence.container, sequence.prefix),
            name_prefix = sequence.prefix,
            suffix = sequence.suffix,
            first = sequence.first.number,
            last = sequence.last.number,
            field = "#" * sequence.first.digit_count,
        )
    return generate

def format_ambiguity(amb):
    template = (
        MSG_AMBIGUOUS_PREVIOUS
        if amb.direction == Collector.AMBIGUITY_PREVIOUS
        else MSG_AMBIGUOUS_NEXT
    )
    return template.format(**amb._as_dict())

def group(iterable, key=None):
    if key is None:
        key = lambda m: m

    result = collections.OrderedDict()

    for m in iterable:
        k = key(m)
        if k not in result:
            result[k] = [ m ]
        else:
            result[k].append(m)

    for k, ms in result.items():
        yield k, ms
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright (C) 2003-2009 Edgewall Software
# Copyright (C) 2003-2004 Jonas Borgström <jonas@edgewall.com>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.org/wiki/TracLicense.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://trac.edgewall.org/log/.
#
# Author: Jonas Borgström <jonas@edgewall.com>

import pkg_resources
import re

from genshi.builder import tag, Element

from trac.config import IntOption, ListOption
from trac.core import *
from trac.mimeview import Context
from trac.perm import IPermissionRequestor
from trac.search.api import ISearchSource
from trac.util.datefmt import format_datetime
from trac.util.presentation import Paginator
from trac.util.translation import _
from trac.web import IRequestHandler
from trac.web.chrome import add_link, add_stylesheet, add_warning, \
                            INavigationContributor, ITemplateProvider
from trac.wiki.api import IWikiSyntaxProvider
from trac.wiki.formatter import extract_link

class SearchModule(Component):

    implements(INavigationContributor, IPermissionRequestor, IRequestHandler,
               ITemplateProvider, IWikiSyntaxProvider)

    search_sources = ExtensionPoint(ISearchSource)

    RESULTS_PER_PAGE = 10

    min_query_length = IntOption('search', 'min_query_length', 3,
        """Minimum length of query string allowed when performing a search.""")

    default_disabled_filters = ListOption('search', 'default_disabled_filters',
        doc="""Specifies which search filters should be disabled by default
               on the search page. This will also restrict the filters for the
               quick search function. Search filters can still be manually
               enabled by the user on the search page. (since 0.12)""")

    # INavigationContributor methods

    def get_active_navigation_item(self, req):
        return 'search'

    def get_navigation_items(self, req):
        if 'SEARCH_VIEW' in req.perm:
            yield ('mainnav', 'search',
                   tag.a(_('Search'), href=req.href.search(), accesskey=4))

    # IPermissionRequestor methods

    def get_permission_actions(self):
        return ['SEARCH_VIEW']

    # IRequestHandler methods

    def match_request(self, req):
        return re.match(r'/search(?:/opensearch)?$', req.path_info) is not None

    def process_request(self, req):
        req.perm.assert_permission('SEARCH_VIEW')

        if req.path_info == '/search/opensearch':
            return ('opensearch.xml', {},
                    'application/opensearchdescription+xml')

        query = req.args.get('q')
        available_filters = []
        for source in self.search_sources:
            available_filters.extend(source.get_search_filters(req) or [])

        filters = self._get_selected_filters(req, available_filters)
        data = self._prepare_data(req, query, available_filters, filters)
        if query:
            data['quickjump'] = self._check_quickjump(req, query)
            if query.startswith('!'):
                query = query[1:]

            terms = self._parse_query(req, query)
            if terms:
                results = self._do_search(req, terms, filters)
                if results:
                    data.update(self._prepare_results(req, filters, results))

        add_stylesheet(req, 'common/css/search.css')
        return 'search.html', data, None

    # ITemplateProvider methods

    def get_htdocs_dirs(self):
        return []

    def get_templates_dirs(self):
        return [pkg_resources.resource_filename('trac.search', 'templates')]

    # IWikiSyntaxProvider methods

    def get_wiki_syntax(self):
        return []

    def get_link_resolvers(self):
        yield ('search', self._format_link)

    def _format_link(self, formatter, ns, target, label):
        path, query, fragment = formatter.split_link(target)
        if query:
            href = formatter.href.search() + query.replace(' ', '+')
        else:
            href = formatter.href.search(q=path)
        return tag.a(label, class_='search', href=href)

    # IRequestHandler helper methods

    def _get_selected_filters(self, req, available_filters):
        """Return selected filters or the default filters if none was selected.
        """
        filters = [f[0] for f in available_filters if f[0] in req.args]
        if not filters:
            filters = [f[0] for f in available_filters
                       if f[0] not in self.default_disabled_filters and
                       (len(f) < 3 or len(f) > 2 and f[2])]
        return filters

    def _prepare_data(self, req, query, available_filters, filters):
        return {'filters': [{'name': f[0], 'label': f[1],
                             'active': f[0] in filters}
                            for f in available_filters],
                'query': query, 'quickjump': None, 'results': []}

    def _check_quickjump(self, req, kwd):
        """Look for search shortcuts"""
        noquickjump = int(req.args.get('noquickjump', '0'))
        # Source quickjump   FIXME: delegate to ISearchSource.search_quickjump
        quickjump_href = None
        if kwd[0] == '/':
            quickjump_href = req.href.browser(kwd)
            name = kwd
            description = _('Browse repository path %(path)s', path=kwd)
        else:
            link = extract_link(self.env, Context.from_request(req, 'search'),
                                kwd)
            if isinstance(link, Element):
                quickjump_href = link.attrib.get('href')
                name = link.children
                description = link.attrib.get('title', '')
        if quickjump_href:
            # Only automatically redirect to local quickjump links
            if not quickjump_href.startswith(req.base_path or '/'):
                noquickjump = True
            if noquickjump:
                return {'href': quickjump_href, 'name': tag.EM(name),
                        'description': description}
            else:
                req.redirect(quickjump_href)

    def _get_search_terms(self, query):
        """Break apart a search query into its various search terms.

        Terms are grouped implicitly by word boundary, or explicitly by (single
        or double) quotes.
        """
        terms = []
        for term in re.split('(".*?")|(\'.*?\')|(\s+)', query):
            if term is not None and term.strip():
                if term[0] == term[-1] and term[0] in "'\"":
                    term = term[1:-1]
                terms.append(term)
        return terms

    def _parse_query(self, req, query):
        """Parse query and refuse those which would result in a huge result set
        """
        terms = self._get_search_terms(query)
        if terms and (len(terms) > 1 or
                      len(terms[0]) >= self.min_query_length):
            return terms

        add_warning(req, _('Search query too short. '
                           'Query must be at least %(num)s characters long.',
                           num=self.min_query_length))

    def _do_search(self, req, terms, filters):
        results = []
        for source in self.search_sources:
            results.extend(source.get_search_results(req, terms, filters)
                           or [])
        return sorted(results, key=lambda x: x[2], reverse=True)

    def _prepare_results(self, req, filters, results):
        page = int(req.args.get('page', '1'))
        results = Paginator(results, page - 1, self.RESULTS_PER_PAGE)
        for idx, result in enumerate(results):
            results[idx] = {'href': result[0], 'title': result[1],
                            'date': format_datetime(result[2]),
                            'author': result[3], 'excerpt': result[4]}

        pagedata = []
        shown_pages = results.get_shown_pages(21)
        for shown_page in shown_pages:
            page_href = req.href.search([(f, 'on') for f in filters],
                                        q=req.args.get('q'),
                                        page=shown_page, noquickjump=1)
            pagedata.append([page_href, None, str(shown_page),
                             'page ' + str(shown_page)])

        fields = ['href', 'class', 'string', 'title']
        results.shown_pages = [dict(zip(fields, p)) for p in pagedata]

        results.current_page = {'href': None, 'class': 'current',
                                'string': str(results.page + 1),
                                'title':None}

        if results.has_next_page:
            next_href = req.href.search(zip(filters, ['on'] * len(filters)),
                                        q=req.args.get('q'), page=page + 1,
                                        noquickjump=1)
            add_link(req, 'next', next_href, _('Next Page'))

        if results.has_previous_page:
            prev_href = req.href.search(zip(filters, ['on'] * len(filters)),
                                        q=req.args.get('q'), page=page - 1,
                                        noquickjump=1)
            add_link(req, 'prev', prev_href, _('Previous Page'))

        page_href = req.href.search(
            zip(filters, ['on'] * len(filters)), q=req.args.get('q'),
            noquickjump=1)
        return {'results': results, 'page_href': page_href}
<commit_msg>0.12.2dev: improve documentation for `[search] default_disabled_filters`.<commit_after># -*- coding: utf-8 -*-
#
# Copyright (C) 2003-2009 Edgewall Software
# Copyright (C) 2003-2004 Jonas Borgström <jonas@edgewall.com>
# All rights reserved.
#
# This software is licensed as described in the file COPYING, which
# you should have received as part of this distribution. The terms
# are also available at http://trac.edgewall.org/wiki/TracLicense.
#
# This software consists of voluntary contributions made by many
# individuals. For the exact contribution history, see the revision
# history and logs, available at http://trac.edgewall.org/log/.
#
# Author: Jonas Borgström <jonas@edgewall.com>

import pkg_resources
import re

from genshi.builder import tag, Element

from trac.config import IntOption, ListOption
from trac.core import *
from trac.mimeview import Context
from trac.perm import IPermissionRequestor
from trac.search.api import ISearchSource
from trac.util.datefmt import format_datetime
from trac.util.presentation import Paginator
from trac.util.translation import _
from trac.web import IRequestHandler
from trac.web.chrome import add_link, add_stylesheet, add_warning, \
                            INavigationContributor, ITemplateProvider
from trac.wiki.api import IWikiSyntaxProvider
from trac.wiki.formatter import extract_link

class SearchModule(Component):

    implements(INavigationContributor, IPermissionRequestor, IRequestHandler,
               ITemplateProvider, IWikiSyntaxProvider)

    search_sources = ExtensionPoint(ISearchSource)

    RESULTS_PER_PAGE = 10

    min_query_length = IntOption('search', 'min_query_length', 3,
        """Minimum length of query string allowed when performing a search.""")

    default_disabled_filters = ListOption('search', 'default_disabled_filters',
        doc="""Specifies which search filters should be disabled by
               default on the search page. This will also restrict the
               filters for the quick search function. The filter names
               defined by default components are: `wiki`, `ticket`,
               `milestone` and `changeset`.  For plugins, look for
               their implementation of the ISearchSource interface, in
               the `get_search_filters()` method, the first member of
               returned tuple. Once disabled, search filters can still
               be manually enabled by the user on the search page.
               (since 0.12)""")

    # INavigationContributor methods

    def get_active_navigation_item(self, req):
        return 'search'

    def get_navigation_items(self, req):
        if 'SEARCH_VIEW' in req.perm:
            yield ('mainnav', 'search',
                   tag.a(_('Search'), href=req.href.search(), accesskey=4))

    # IPermissionRequestor methods

    def get_permission_actions(self):
        return ['SEARCH_VIEW']

    # IRequestHandler methods

    def match_request(self, req):
        return re.match(r'/search(?:/opensearch)?$', req.path_info) is not None

    def process_request(self, req):
        req.perm.assert_permission('SEARCH_VIEW')

        if req.path_info == '/search/opensearch':
            return ('opensearch.xml', {},
                    'application/opensearchdescription+xml')

        query = req.args.get('q')
        available_filters = []
        for source in self.search_sources:
            available_filters.extend(source.get_search_filters(req) or [])

        filters = self._get_selected_filters(req, available_filters)
        data = self._prepare_data(req, query, available_filters, filters)
        if query:
            data['quickjump'] = self._check_quickjump(req, query)
            if query.startswith('!'):
                query = query[1:]

            terms = self._parse_query(req, query)
            if terms:
                results = self._do_search(req, terms, filters)
                if results:
                    data.update(self._prepare_results(req, filters, results))

        add_stylesheet(req, 'common/css/search.css')
        return 'search.html', data, None

    # ITemplateProvider methods

    def get_htdocs_dirs(self):
        return []

    def get_templates_dirs(self):
        return [pkg_resources.resource_filename('trac.search', 'templates')]

    # IWikiSyntaxProvider methods

    def get_wiki_syntax(self):
        return []

    def get_link_resolvers(self):
        yield ('search', self._format_link)

    def _format_link(self, formatter, ns, target, label):
        path, query, fragment = formatter.split_link(target)
        if query:
            href = formatter.href.search() + query.replace(' ', '+')
        else:
            href = formatter.href.search(q=path)
        return tag.a(label, class_='search', href=href)

    # IRequestHandler helper methods

    def _get_selected_filters(self, req, available_filters):
        """Return selected filters or the default filters if none was selected.
        """
        filters = [f[0] for f in available_filters if f[0] in req.args]
        if not filters:
            filters = [f[0] for f in available_filters
                       if f[0] not in self.default_disabled_filters and
                       (len(f) < 3 or len(f) > 2 and f[2])]
        return filters

    def _prepare_data(self, req, query, available_filters, filters):
        return {'filters': [{'name': f[0], 'label': f[1],
                             'active': f[0] in filters}
                            for f in available_filters],
                'query': query, 'quickjump': None, 'results': []}

    def _check_quickjump(self, req, kwd):
        """Look for search shortcuts"""
        noquickjump = int(req.args.get('noquickjump', '0'))
        # Source quickjump   FIXME: delegate to ISearchSource.search_quickjump
        quickjump_href = None
        if kwd[0] == '/':
            quickjump_href = req.href.browser(kwd)
            name = kwd
            description = _('Browse repository path %(path)s', path=kwd)
        else:
            link = extract_link(self.env, Context.from_request(req, 'search'),
                                kwd)
            if isinstance(link, Element):
                quickjump_href = link.attrib.get('href')
                name = link.children
                description = link.attrib.get('title', '')
        if quickjump_href:
            # Only automatically redirect to local quickjump links
            if not quickjump_href.startswith(req.base_path or '/'):
                noquickjump = True
            if noquickjump:
                return {'href': quickjump_href, 'name': tag.EM(name),
                        'description': description}
            else:
                req.redirect(quickjump_href)

    def _get_search_terms(self, query):
        """Break apart a search query into its various search terms.

        Terms are grouped implicitly by word boundary, or explicitly by (single
        or double) quotes.
        """
        terms = []
        for term in re.split('(".*?")|(\'.*?\')|(\s+)', query):
            if term is not None and term.strip():
                if term[0] == term[-1] and term[0] in "'\"":
                    term = term[1:-1]
                terms.append(term)
        return terms

    def _parse_query(self, req, query):
        """Parse query and refuse those which would result in a huge result set
        """
        terms = self._get_search_terms(query)
        if terms and (len(terms) > 1 or
                      len(terms[0]) >= self.min_query_length):
            return terms

        add_warning(req, _('Search query too short. '
                           'Query must be at least %(num)s characters long.',
                           num=self.min_query_length))

    def _do_search(self, req, terms, filters):
        results = []
        for source in self.search_sources:
            results.extend(source.get_search_results(req, terms, filters)
                           or [])
        return sorted(results, key=lambda x: x[2], reverse=True)

    def _prepare_results(self, req, filters, results):
        page = int(req.args.get('page', '1'))
        results = Paginator(results, page - 1, self.RESULTS_PER_PAGE)
        for idx, result in enumerate(results):
            results[idx] = {'href': result[0], 'title': result[1],
                            'date': format_datetime(result[2]),
                            'author': result[3], 'excerpt': result[4]}

        pagedata = []
        shown_pages = results.get_shown_pages(21)
        for shown_page in shown_pages:
            page_href = req.href.search([(f, 'on') for f in filters],
                                        q=req.args.get('q'),
                                        page=shown_page, noquickjump=1)
            pagedata.append([page_href, None, str(shown_page),
                             'page ' + str(shown_page)])

        fields = ['href', 'class', 'string', 'title']
        results.shown_pages = [dict(zip(fields, p)) for p in pagedata]

        results.current_page = {'href': None, 'class': 'current',
                                'string': str(results.page + 1),
                                'title':None}

        if results.has_next_page:
            next_href = req.href.search(zip(filters, ['on'] * len(filters)),
                                        q=req.args.get('q'), page=page + 1,
                                        noquickjump=1)
            add_link(req, 'next', next_href, _('Next Page'))

        if results.has_previous_page:
            prev_href = req.href.search(zip(filters, ['on'] * len(filters)),
                                        q=req.args.get('q'), page=page - 1,
                                        noquickjump=1)
            add_link(req, 'prev', prev_href, _('Previous Page'))

        page_href = req.href.search(
            zip(filters, ['on'] * len(filters)), q=req.args.get('q'),
            noquickjump=1)
        return {'results': results, 'page_href': page_href}
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# -*- mode:python; sh-basic-offset:4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim:set tabstop=4 softtabstop=4 expandtab shiftwidth=4 fileencoding=utf-8:
#

__version__ = '0.9.1'
<commit_msg>Bumping version for release<commit_after>#!/usr/bin/env python
#
# -*- mode:python; sh-basic-offset:4; indent-tabs-mode:nil; coding:utf-8 -*-
# vim:set tabstop=4 softtabstop=4 expandtab shiftwidth=4 fileencoding=utf-8:
#

__version__ = '0.9.2'
<|endoftext|>