<commit_before>""" Miscallanious util funcitons for the mix-mind project
"""
from functools import wraps
from fractions import Fraction
from collections import OrderedDict, namedtuple
import operator
import json
import inspect
import uuid
import pendulum
from . import log

# make passing a bunch of options around a bit cleaner
DisplayOptions = namedtuple('DisplayOptions', 'prices,stats,examples,all_ingredients,markup,prep_line,origin,info,variants')
FilterOptions = namedtuple('FilterOptions', 'search,all_,include,exclude,include_use_or,exclude_use_or,style,glass,prep,ice,name,tag')
PdfOptions = namedtuple('PdfOptions', 'pdf_filename,ncols,liquor_list,liquor_list_own_page,debug,align,title,tagline')

VALID_UNITS = ['oz', 'mL', 'cL']

class ResultRecipes(object):
    def add_items(self, recipes):
        pass
    def get_items(self):
        pass

class UnionResultRecipes(ResultRecipes):
    def __init__(self):
        self.container = OrderedDict()
    def add_items(self, recipes):
        for recipe in recipes:
            self.container[recipe.name] = recipe
    def get_items(self):
        return self.container.values()

class IntersectionResultRecipes(ResultRecipes):
    def __init__(self):
        self.container = None
    def add_items(self, recipes):
        if self.container is None:
            self.container = recipes
        else:
            self.container = [x for x in self.container if x in recipes]
    def get_items(self):
        return self.container

def filter_recipes(all_recipes, filter_options, union_results=False):
    """Filters the recipe list based on a FilterOptions bundle of parameters
    :param list[Recipe] all_recipes: list of recipe object to filter
    :param FilterOptions filter_options: bundle of filtering parameters
        search str: search an arbitrary string against the ingredients and attributes
    :param bool union_results: for each attributes searched against, combine results
        with set intersection by default, or union if True
    """
    result_recipes = UnionResultRecipes() if union_results else IntersectionResultRecipes()
    recipes = [recipe for recipe in all_recipes if filter_options.all_ or recipe.can_make]
    if filter_options.search:
        include_list = [filter_options.search.lower()]
    else:
        include_list = filter_options.include
    if include_list:
        reduce_fn = any if filter_options.include_use_or else all
        result_recipes.add_items([recipe for recipe in recipes if
                reduce_fn((recipe.contains_ingredient(ingredient, include_optional=True)
                for ingredient in include_list))])
    if filter_options.exclude:
        reduce_fn = any if filter_options.exclude_use_or else all
        result_recipes.add_items([recipe for recipe in recipes if
                reduce_fn((not recipe.contains_ingredient(ingredient, include_optional=False)
                for ingredient in filter_options.exclude))])
    for attr in 'style glass prep ice name tag'.split():
        result_recipes.add_items(filter_on_attribute(recipes, filter_options, attr))

    result_recipes = result_recipes.get_items()

    def get_names(items):
        return set(map(lambda i: i.name, items))
    excluded = sorted(list(get_names(all_recipes) - get_names(result_recipes)))
    log.debug("Excluded: {}\n".format(', '.join(excluded)))
    return result_recipes, excluded

def find_recipe(recipes, name):
    for recipe in recipes:
        if recipe.name == name:
            return recipe
    return None

def filter_on_attribute(recipes, filter_options, attribute):
    attr_value = getattr(filter_options, attribute).lower()
    if filter_options.search and not attr_value:
        attr_value = filter_options.search.lower()
    if attr_value:
        recipes = [recipe for recipe in recipes if attr_value in getattr(recipe, attribute).lower()]
    return recipes

def get_uuid():
    return unicode(uuid.uuid4())

class StatTracker(dict):
    # mutable class variables
    _title_width = 0
    _name_width = 0

    def __init__(self, attr, magnitude, str_title):
        if magnitude not in ('max', 'min'):
            raise ValueError('StatTracker magnitude must be "max" or "min"')
        self.op = operator.lt if magnitude == 'min' else operator.gt
        self.stat = '{}_{}'.format(magnitude, attr)
        self.val_attr = attr
        self.val = float('inf') if magnitude == 'min' else 0.0
        self['title'] = str_title
        if len(str_title) > StatTracker._title_width:
            StatTracker._title_width = len(str_title)

    def __str__(self):
        return "{{title:{}}} | {{drink_name:{}}} | ${{cost:.2f}} | {{abv:>5.2f}}% ABV | {{std_drinks:.2f}} | {{bottles}}"\
            .format(self._title_width+1, self._name_width+1).format(**self)

    def as_html(self):
        return u"
 {{title:{}}} 	 {{drink_name:{}}} 	 ${{cost:.2f}} 	 {{abv:>5.2f}}% ABV 	 {{std_drinks:.2f}} 	 {{bottles}}
"\
            .format(self._title_width+1, self._name_width+1).format(**self)

    def update_stat(self, recipe):
        example = getattr(recipe.stats, self.stat)
        ex_val = getattr(example, self.val_attr)
        if self.op(ex_val, self.val):
            self.val = ex_val
            self.update(example._asdict())
            self['drink_name'] = recipe.name
            if len(recipe.name) > StatTracker._name_width:
                StatTracker._name_width = len(recipe.name)

def report_stats(recipes, as_html=False):
    most_expensive = StatTracker('cost', 'max', 'Most Expensive')
    most_booze = StatTracker('std_drinks', 'max', 'Most Std Drinks')
    most_abv = StatTracker('abv', 'max', 'Highest Estimated ABV')
    least_expensive = StatTracker('cost', 'min', 'Least Expensive')
    least_booze = StatTracker('std_drinks', 'min', 'Fewest Std Drinks')
    least_abv = StatTracker('abv', 'min', 'Lowest Estimated ABV')
    for recipe in recipes:
        if recipe.calculate_stats():
            most_expensive.update_stat(recipe)
            most_booze.update_stat(recipe)
            most_abv.update_stat(recipe)
            least_expensive.update_stat(recipe)
            least_booze.update_stat(recipe)
            least_abv.update_stat(recipe)
    if as_html:
        return u"

	Drink	Cost	Est ABV	Std Drinks	Ingredients

{}

".format(u''.join([s.as_html()
            for s in [most_expensive, most_booze, most_abv, least_expensive, least_booze, least_abv]]))
    else:
        return [most_expensive, most_booze, most_abv, least_expensive, least_booze, least_abv]

def load_recipe_json(recipe_files):
    base_recipes = OrderedDict()
    for recipe_json in recipe_files:
        with open(recipe_json) as fp:
            other_recipes = json.load(fp, object_pairs_hook=OrderedDict)
            log.info("Recipes loaded from {}".format(recipe_json))
            for item in other_recipes.itervalues():
                item.update({'source_file': recipe_json})
            for name in [name for name in other_recipes.keys() if name in base_recipes.keys()]:
                log.debug("Keeping {} from {} over {}".format(name, base_recipes[name]['source_file'], other_recipes[name]['source_file']))
                del other_recipes[name]
            base_recipes.update(other_recipes)
    return base_recipes

def default_initializer(func):
    names, varargs, keywords, defaults = inspect.getargspec(func)
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        for name, arg in list(zip(names[1:], args)) + list(kwargs.items()):
            setattr(self, name, arg)
        func(self, *args, **kwargs)
    return wrapper

# utils to convert string values
def from_float(s):
    if not s:
        return 0.0
    return float(s)
def from_price_float(s):
    if isinstance(s, basestring):
        return from_float(s.replace('$', ''))
    return from_float(s)
def from_bool_from_num(s):
    if not s:
        return False
    return bool(float(s))
def as_utf8(s):
    return unicode(s, 'utf-8')

def to_fraction(amount):
    fraction = Fraction.from_float(float(amount)).limit_denominator(99)
    if fraction.denominator == 1:
        return fraction.numerator
    whole = fraction.numerator / fraction.denominator
    numer = fraction.numerator % fraction.denominator
    return "{}{}/{}".format(str(whole)+' ' if whole > 0 else '', numer, fraction.denominator)

def calculate_price(cost, markup):
    return int(((cost + 1) * float(markup)) +1)

def calculate_std_drinks(abv, amount, unit):
    """ Standard drink is 1.5 oz or 45 ml at 40% abv
    """
    adjusted_abv = abv / 40.0
    adjusted_amount = convert_units(amount, unit, 'oz') / 1.5
    return adjusted_abv * adjusted_amount

def calculate_abv(std_drinks, volume, unit):
    if unit == 'oz':
        units_per_std_drink = 1.5
    elif unit == 'mL':
        units_per_std_drink = 45.0
    elif unit == 'cL':
        units_per_std_drink = 4.5
    else:
        raise NotImplementedError("number of standard drinks for unit '{}' is unknown".format(unit))
    abv = 40.0 * (std_drinks*units_per_std_drink / volume)
    return abv

# units, yo
ML_PER_OZ           =  29.5735
ML_PER_OZ_ROUNDED   =  30.0
ML_PER_TSP          =  4.92892
ML_PER_TSP_ROUNDED  =  5.0
ML_PER_DS           =  0.92
ML_PER_CL           =  10.0
ML_PER_DROP         =  0.12
OZ_PER_TSP          =  1.0/8.0
OZ_PER_DS           =  1.0/32.0
OZ_PER_DROP         =  1.0/240.0

def convert_units(amount, from_unit, to_unit, rounded=False):
    if from_unit == 'literal':
        return amount
    try:
        amount = float(amount)
    except TypeError: # pd series breaks this
        amount = amount
    if from_unit == to_unit:
        return amount
    unit_conversions = {
            'ds': dash_to_volume,
            'tsp': tsp_to_volume,
            'mL': mL_to_volume,
            'cL': cL_to_volume,
            'oz': oz_to_volume,
            'drop': drop_to_volume,
            }
    convert = unit_conversions.get(from_unit, lambda x,y,z: no_conversion(from_unit, to_unit))
    return convert(amount, to_unit, rounded)

def no_conversion(from_unit, to_unit):
    raise NotImplementedError("conversion from {} to {}".format(from_unit, to_unit))

def dash_to_volume(amount, unit, rounded=False):
    mL_per_oz = ML_PER_OZ if not rounded else ML_PER_OZ_ROUNDED
    if unit == 'mL':
        return amount * ML_PER_DS
    elif unit == 'cL':
        return amount * ML_PER_DS / ML_PER_CL
    elif unit == 'oz':
        return amount / mL_per_oz
    else:
        no_conversion('dash', unit)

def tsp_to_volume(amount, unit, rounded=False):
    mL_per_tsp = ML_PER_TSP if not rounded else ML_PER_TSP_ROUNDED
    if unit == 'oz':
        return amount * OZ_PER_TSP
    elif unit == 'mL':
        return amount * mL_per_tsp
    elif unit == 'cL':
        return amount * mL_per_tsp / ML_PER_CL
    else:
        no_conversion('tsp', unit)

def oz_to_volume(amount, unit, rounded=False):
    mL_per_oz = ML_PER_OZ if not rounded else ML_PER_OZ_ROUNDED
    if unit == 'mL':
        return amount * mL_per_oz
    elif unit == 'cL':
        return amount * mL_per_oz / ML_PER_CL
    elif unit == 'tsp':
        return amount / OZ_PER_TSP
    elif unit == 'ds':
        return amount / OZ_PER_DS
    elif unit == 'drop':
        return amount / OZ_PER_DROP
    else:
        no_conversion('oz', unit)

def mL_to_volume(amount, unit, rounded=False):
    mL_per_oz = ML_PER_OZ if not rounded else ML_PER_OZ_ROUNDED
    mL_per_tsp = ML_PER_TSP if not rounded else ML_PER_TSP_ROUNDED
    if unit == 'oz':
        return amount / mL_per_oz
    elif unit == 'cL':
        return amount / ML_PER_CL
    elif unit == 'ds':
        return amount / ML_PER_DS
    elif unit == 'tsp':
        return amount / mL_per_tsp
    elif unit == 'drop':
        return amount / ML_PER_DROP
    elif unit == 'mL':
        return amount
    else:
        no_conversion('mL', unit)

def drop_to_volume(amount, unit, rounded=False):
    if unit == 'oz':
        return amount * OZ_PER_DROP
    elif unit == 'mL':
        return amount * ML_PER_DROP
    elif unit == 'cL':
        return amount * ML_PER_DROP / ML_PER_CL
    else:
        no_conversion('drop', unit)

def cL_to_volume(amount, unit, rounded=False):
    try:
        return mL_to_volume(amount, unit, rounded) * ML_PER_CL
    except NotImplementedError:
        no_conversion('cL', unit)

class IngredientSpecifier(object):
    """ Allow type:bottle in recipes,
    e.g. "white rum:Barcadi Catra Blanca" or "aromatic bitters:Angostura"
    """
    @default_initializer
    def __init__(self, what, bottle=None):
        if what is None:
            raise ValueError("IngredientSpecifier what (type) cannot be None")
        if '(' in what and ')' in what:
            self.extra = what.strip()[what.find('('):]
            self.what = what.strip()[:what.find('(')].strip()
        else:
            self.extra = None

    @classmethod
    def from_string(cls, type_str):
        if ':' in type_str:
            t = type_str.split(':')
            if len(t) == 2:
                what = t[0]
                bottle = t[1]
            else:
                raise ValueError("Unknown ingredient specifier: {}".format(type_str))
        else:
            what = type_str
            bottle = None
        return cls(what, bottle)

    def __str__(self):
        return self.bottle if self.bottle else u"{}{}".format(self.what, u' '+self.extra if self.extra else u'')

    def __repr__(self):
        return u"{}:{}".format(self.what, self.bottle if self.bottle else '')

def to_human_diff(dt):
    """Return datetime as humanized diff from now"""
    return pendulum.instance(dt).diff_for_humans() if dt else '-'

def get_ts_formatter(fmt, tz):
    """Returns callable that will format a datetime"""
    return lambda dt: pendulum.instance(dt).in_timezone(tz).format(fmt) if dt else '-'

<commit_msg>Handle unicode decoding unicode error<commit_after>""" Miscallanious util funcitons for the mix-mind project
"""
from functools import wraps
from fractions import Fraction
from collections import OrderedDict, namedtuple
import operator
import json
import inspect
import uuid
import pendulum
from . import log

# make passing a bunch of options around a bit cleaner
DisplayOptions = namedtuple('DisplayOptions', 'prices,stats,examples,all_ingredients,markup,prep_line,origin,info,variants')
FilterOptions = namedtuple('FilterOptions', 'search,all_,include,exclude,include_use_or,exclude_use_or,style,glass,prep,ice,name,tag')
PdfOptions = namedtuple('PdfOptions', 'pdf_filename,ncols,liquor_list,liquor_list_own_page,debug,align,title,tagline')

VALID_UNITS = ['oz', 'mL', 'cL']

class ResultRecipes(object):
    def add_items(self, recipes):
        pass
    def get_items(self):
        pass

class UnionResultRecipes(ResultRecipes):
    def __init__(self):
        self.container = OrderedDict()
    def add_items(self, recipes):
        for recipe in recipes:
            self.container[recipe.name] = recipe
    def get_items(self):
        return self.container.values()

class IntersectionResultRecipes(ResultRecipes):
    def __init__(self):
        self.container = None
    def add_items(self, recipes):
        if self.container is None:
            self.container = recipes
        else:
            self.container = [x for x in self.container if x in recipes]
    def get_items(self):
        return self.container

def filter_recipes(all_recipes, filter_options, union_results=False):
    """Filters the recipe list based on a FilterOptions bundle of parameters
    :param list[Recipe] all_recipes: list of recipe object to filter
    :param FilterOptions filter_options: bundle of filtering parameters
        search str: search an arbitrary string against the ingredients and attributes
    :param bool union_results: for each attributes searched against, combine results
        with set intersection by default, or union if True
    """
    result_recipes = UnionResultRecipes() if union_results else IntersectionResultRecipes()
    recipes = [recipe for recipe in all_recipes if filter_options.all_ or recipe.can_make]
    if filter_options.search:
        include_list = [filter_options.search.lower()]
    else:
        include_list = filter_options.include
    if include_list:
        reduce_fn = any if filter_options.include_use_or else all
        result_recipes.add_items([recipe for recipe in recipes if
                reduce_fn((recipe.contains_ingredient(ingredient, include_optional=True)
                for ingredient in include_list))])
    if filter_options.exclude:
        reduce_fn = any if filter_options.exclude_use_or else all
        result_recipes.add_items([recipe for recipe in recipes if
                reduce_fn((not recipe.contains_ingredient(ingredient, include_optional=False)
                for ingredient in filter_options.exclude))])
    for attr in 'style glass prep ice name tag'.split():
        result_recipes.add_items(filter_on_attribute(recipes, filter_options, attr))

    result_recipes = result_recipes.get_items()

    def get_names(items):
        return set(map(lambda i: i.name, items))
    excluded = sorted(list(get_names(all_recipes) - get_names(result_recipes)))
    log.debug("Excluded: {}\n".format(', '.join(excluded)))
    return result_recipes, excluded

def find_recipe(recipes, name):
    for recipe in recipes:
        if recipe.name == name:
            return recipe
    return None

def filter_on_attribute(recipes, filter_options, attribute):
    attr_value = getattr(filter_options, attribute).lower()
    if filter_options.search and not attr_value:
        attr_value = filter_options.search.lower()
    if attr_value:
        recipes = [recipe for recipe in recipes if attr_value in getattr(recipe, attribute).lower()]
    return recipes

def get_uuid():
    return unicode(uuid.uuid4())

class StatTracker(dict):
    # mutable class variables
    _title_width = 0
    _name_width = 0

    def __init__(self, attr, magnitude, str_title):
        if magnitude not in ('max', 'min'):
            raise ValueError('StatTracker magnitude must be "max" or "min"')
        self.op = operator.lt if magnitude == 'min' else operator.gt
        self.stat = '{}_{}'.format(magnitude, attr)
        self.val_attr = attr
        self.val = float('inf') if magnitude == 'min' else 0.0
        self['title'] = str_title
        if len(str_title) > StatTracker._title_width:
            StatTracker._title_width = len(str_title)

    def __str__(self):
        return "{{title:{}}} | {{drink_name:{}}} | ${{cost:.2f}} | {{abv:>5.2f}}% ABV | {{std_drinks:.2f}} | {{bottles}}"\
            .format(self._title_width+1, self._name_width+1).format(**self)

    def as_html(self):
        return u"
 {{title:{}}} 	 {{drink_name:{}}} 	 ${{cost:.2f}} 	 {{abv:>5.2f}}% ABV 	 {{std_drinks:.2f}} 	 {{bottles}}
"\
            .format(self._title_width+1, self._name_width+1).format(**self)

    def update_stat(self, recipe):
        example = getattr(recipe.stats, self.stat)
        ex_val = getattr(example, self.val_attr)
        if self.op(ex_val, self.val):
            self.val = ex_val
            self.update(example._asdict())
            self['drink_name'] = recipe.name
            if len(recipe.name) > StatTracker._name_width:
                StatTracker._name_width = len(recipe.name)

def report_stats(recipes, as_html=False):
    most_expensive = StatTracker('cost', 'max', 'Most Expensive')
    most_booze = StatTracker('std_drinks', 'max', 'Most Std Drinks')
    most_abv = StatTracker('abv', 'max', 'Highest Estimated ABV')
    least_expensive = StatTracker('cost', 'min', 'Least Expensive')
    least_booze = StatTracker('std_drinks', 'min', 'Fewest Std Drinks')
    least_abv = StatTracker('abv', 'min', 'Lowest Estimated ABV')
    for recipe in recipes:
        if recipe.calculate_stats():
            most_expensive.update_stat(recipe)
            most_booze.update_stat(recipe)
            most_abv.update_stat(recipe)
            least_expensive.update_stat(recipe)
            least_booze.update_stat(recipe)
            least_abv.update_stat(recipe)
    if as_html:
        return u"

	Drink	Cost	Est ABV	Std Drinks	Ingredients

{}

".format(u''.join([s.as_html()
            for s in [most_expensive, most_booze, most_abv, least_expensive, least_booze, least_abv]]))
    else:
        return [most_expensive, most_booze, most_abv, least_expensive, least_booze, least_abv]

def load_recipe_json(recipe_files):
    base_recipes = OrderedDict()
    for recipe_json in recipe_files:
        with open(recipe_json) as fp:
            other_recipes = json.load(fp, object_pairs_hook=OrderedDict)
            log.info("Recipes loaded from {}".format(recipe_json))
            for item in other_recipes.itervalues():
                item.update({'source_file': recipe_json})
            for name in [name for name in other_recipes.keys() if name in base_recipes.keys()]:
                log.debug("Keeping {} from {} over {}".format(name, base_recipes[name]['source_file'], other_recipes[name]['source_file']))
                del other_recipes[name]
            base_recipes.update(other_recipes)
    return base_recipes

def default_initializer(func):
    names, varargs, keywords, defaults = inspect.getargspec(func)
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        for name, arg in list(zip(names[1:], args)) + list(kwargs.items()):
            setattr(self, name, arg)
        func(self, *args, **kwargs)
    return wrapper

# utils to convert string values
def from_float(s):
    if not s:
        return 0.0
    return float(s)
def from_price_float(s):
    if isinstance(s, basestring):
        return from_float(s.replace('$', ''))
    return from_float(s)
def from_bool_from_num(s):
    if not s:
        return False
    return bool(float(s))
def as_utf8(s):
    try:
        return unicode(s, 'utf-8')
    except TypeError as e:
        if "decoding Unicode is not supported" in str(e):
            return s
        raise

def to_fraction(amount):
    fraction = Fraction.from_float(float(amount)).limit_denominator(99)
    if fraction.denominator == 1:
        return fraction.numerator
    whole = fraction.numerator / fraction.denominator
    numer = fraction.numerator % fraction.denominator
    return "{}{}/{}".format(str(whole)+' ' if whole > 0 else '', numer, fraction.denominator)

def calculate_price(cost, markup):
    return int(((cost + 1) * float(markup)) +1)

def calculate_std_drinks(abv, amount, unit):
    """ Standard drink is 1.5 oz or 45 ml at 40% abv
    """
    adjusted_abv = abv / 40.0
    adjusted_amount = convert_units(amount, unit, 'oz') / 1.5
    return adjusted_abv * adjusted_amount

def calculate_abv(std_drinks, volume, unit):
    if unit == 'oz':
        units_per_std_drink = 1.5
    elif unit == 'mL':
        units_per_std_drink = 45.0
    elif unit == 'cL':
        units_per_std_drink = 4.5
    else:
        raise NotImplementedError("number of standard drinks for unit '{}' is unknown".format(unit))
    abv = 40.0 * (std_drinks*units_per_std_drink / volume)
    return abv

# units, yo
ML_PER_OZ           =  29.5735
ML_PER_OZ_ROUNDED   =  30.0
ML_PER_TSP          =  4.92892
ML_PER_TSP_ROUNDED  =  5.0
ML_PER_DS           =  0.92
ML_PER_CL           =  10.0
ML_PER_DROP         =  0.12
OZ_PER_TSP          =  1.0/8.0
OZ_PER_DS           =  1.0/32.0
OZ_PER_DROP         =  1.0/240.0

def convert_units(amount, from_unit, to_unit, rounded=False):
    if from_unit == 'literal':
        return amount
    try:
        amount = float(amount)
    except TypeError: # pd series breaks this
        amount = amount
    if from_unit == to_unit:
        return amount
    unit_conversions = {
            'ds': dash_to_volume,
            'tsp': tsp_to_volume,
            'mL': mL_to_volume,
            'cL': cL_to_volume,
            'oz': oz_to_volume,
            'drop': drop_to_volume,
            }
    convert = unit_conversions.get(from_unit, lambda x,y,z: no_conversion(from_unit, to_unit))
    return convert(amount, to_unit, rounded)

def no_conversion(from_unit, to_unit):
    raise NotImplementedError("conversion from {} to {}".format(from_unit, to_unit))

def dash_to_volume(amount, unit, rounded=False):
    mL_per_oz = ML_PER_OZ if not rounded else ML_PER_OZ_ROUNDED
    if unit == 'mL':
        return amount * ML_PER_DS
    elif unit == 'cL':
        return amount * ML_PER_DS / ML_PER_CL
    elif unit == 'oz':
        return amount / mL_per_oz
    else:
        no_conversion('dash', unit)

def tsp_to_volume(amount, unit, rounded=False):
    mL_per_tsp = ML_PER_TSP if not rounded else ML_PER_TSP_ROUNDED
    if unit == 'oz':
        return amount * OZ_PER_TSP
    elif unit == 'mL':
        return amount * mL_per_tsp
    elif unit == 'cL':
        return amount * mL_per_tsp / ML_PER_CL
    else:
        no_conversion('tsp', unit)

def oz_to_volume(amount, unit, rounded=False):
    mL_per_oz = ML_PER_OZ if not rounded else ML_PER_OZ_ROUNDED
    if unit == 'mL':
        return amount * mL_per_oz
    elif unit == 'cL':
        return amount * mL_per_oz / ML_PER_CL
    elif unit == 'tsp':
        return amount / OZ_PER_TSP
    elif unit == 'ds':
        return amount / OZ_PER_DS
    elif unit == 'drop':
        return amount / OZ_PER_DROP
    else:
        no_conversion('oz', unit)

def mL_to_volume(amount, unit, rounded=False):
    mL_per_oz = ML_PER_OZ if not rounded else ML_PER_OZ_ROUNDED
    mL_per_tsp = ML_PER_TSP if not rounded else ML_PER_TSP_ROUNDED
    if unit == 'oz':
        return amount / mL_per_oz
    elif unit == 'cL':
        return amount / ML_PER_CL
    elif unit == 'ds':
        return amount / ML_PER_DS
    elif unit == 'tsp':
        return amount / mL_per_tsp
    elif unit == 'drop':
        return amount / ML_PER_DROP
    elif unit == 'mL':
        return amount
    else:
        no_conversion('mL', unit)

def drop_to_volume(amount, unit, rounded=False):
    if unit == 'oz':
        return amount * OZ_PER_DROP
    elif unit == 'mL':
        return amount * ML_PER_DROP
    elif unit == 'cL':
        return amount * ML_PER_DROP / ML_PER_CL
    else:
        no_conversion('drop', unit)

def cL_to_volume(amount, unit, rounded=False):
    try:
        return mL_to_volume(amount, unit, rounded) * ML_PER_CL
    except NotImplementedError:
        no_conversion('cL', unit)

class IngredientSpecifier(object):
    """ Allow type:bottle in recipes,
    e.g. "white rum:Barcadi Catra Blanca" or "aromatic bitters:Angostura"
    """
    @default_initializer
    def __init__(self, what, bottle=None):
        if what is None:
            raise ValueError("IngredientSpecifier what (type) cannot be None")
        if '(' in what and ')' in what:
            self.extra = what.strip()[what.find('('):]
            self.what = what.strip()[:what.find('(')].strip()
        else:
            self.extra = None

    @classmethod
    def from_string(cls, type_str):
        if ':' in type_str:
            t = type_str.split(':')
            if len(t) == 2:
                what = t[0]
                bottle = t[1]
            else:
                raise ValueError("Unknown ingredient specifier: {}".format(type_str))
        else:
            what = type_str
            bottle = None
        return cls(what, bottle)

    def __str__(self):
        return self.bottle if self.bottle else u"{}{}".format(self.what, u' '+self.extra if self.extra else u'')

    def __repr__(self):
        return u"{}:{}".format(self.what, self.bottle if self.bottle else '')

def to_human_diff(dt):
    """Return datetime as humanized diff from now"""
    return pendulum.instance(dt).diff_for_humans() if dt else '-'

def get_ts_formatter(fmt, tz):
    """Returns callable that will format a datetime"""
    return lambda dt: pendulum.instance(dt).in_timezone(tz).format(fmt) if dt else '-'

<|endoftext|>


<commit_before>from __future__ import unicode_literals
import logging
import shutil
import tempfile

from mkdocs.build import build
from mkdocs.config import load_config

log = logging.getLogger(__name__)

def _livereload(host, port, config, builder, site_dir):

    # We are importing here for anyone that has issues with livereload. Even if
    # this fails, the --no-livereload alternative should still work.
    from livereload import Server

    server = Server()

    # Watch the documentation files, the config file and the theme files.
    server.watch(config['docs_dir'], builder)
    server.watch(config['config_file_path'], builder)

    for d in config['theme_dir']:
        server.watch(d, builder)

    server.serve(root=site_dir, host=host, port=int(port), restart_delay=0)

def _static_server(host, port, site_dir):

    # Importing here to seperate the code paths from the --livereload
    # alternative.
    from tornado import ioloop
    from tornado import web

    application = web.Application([
        (r"/(.*)", web.StaticFileHandler, {
            "path": site_dir,
            "default_filename": "index.html"
        }),
    ])
    application.listen(port=port, address=host)

    log.info('Running at: http://%s:%s/', host, port)
    log.info('Hold ctrl+c to quit.')
    try:
        ioloop.IOLoop.instance().start()
    except KeyboardInterrupt:
        log.info('Stopping server...')

def serve(config_file=None, dev_addr=None, strict=None, theme=None,
          livereload=True):
    """
    Start the MkDocs development server

    By default it will serve the documentation on http://localhost:8000/ and
    it will rebuild the documentation and refresh the page automatically
    whenever a file is edited.
    """
    # Create a temporary build directory, and set some options to serve it
    tempdir = tempfile.mkdtemp()

    def builder():
        log.info("Building documentation...")
        config = load_config(
            config_file=config_file,
            dev_addr=dev_addr,
            strict=strict,
            theme=theme,
        )
        config['site_dir'] = tempdir
        build(config, live_server=True)
        return config

    # Perform the initial build
    config = builder()

    host, port = config['dev_addr'].split(':', 1)

    try:
        if livereload:
            _livereload(host, port, config, builder, tempdir)
        else:
            _static_server(host, port, tempdir)
    finally:
        shutil.rmtree(tempdir)
<commit_msg>Perform a clean build when serving<commit_after>from __future__ import unicode_literals
import logging
import shutil
import tempfile

from mkdocs.build import build
from mkdocs.config import load_config

log = logging.getLogger(__name__)

def _livereload(host, port, config, builder, site_dir):

    # We are importing here for anyone that has issues with livereload. Even if
    # this fails, the --no-livereload alternative should still work.
    from livereload import Server

    server = Server()

    # Watch the documentation files, the config file and the theme files.
    server.watch(config['docs_dir'], builder)
    server.watch(config['config_file_path'], builder)

    for d in config['theme_dir']:
        server.watch(d, builder)

    server.serve(root=site_dir, host=host, port=int(port), restart_delay=0)

def _static_server(host, port, site_dir):

    # Importing here to seperate the code paths from the --livereload
    # alternative.
    from tornado import ioloop
    from tornado import web

    application = web.Application([
        (r"/(.*)", web.StaticFileHandler, {
            "path": site_dir,
            "default_filename": "index.html"
        }),
    ])
    application.listen(port=port, address=host)

    log.info('Running at: http://%s:%s/', host, port)
    log.info('Hold ctrl+c to quit.')
    try:
        ioloop.IOLoop.instance().start()
    except KeyboardInterrupt:
        log.info('Stopping server...')

def serve(config_file=None, dev_addr=None, strict=None, theme=None,
          livereload=True):
    """
    Start the MkDocs development server

    By default it will serve the documentation on http://localhost:8000/ and
    it will rebuild the documentation and refresh the page automatically
    whenever a file is edited.
    """
    # Create a temporary build directory, and set some options to serve it
    tempdir = tempfile.mkdtemp()

    def builder():
        log.info("Building documentation...")
        config = load_config(
            config_file=config_file,
            dev_addr=dev_addr,
            strict=strict,
            theme=theme,
        )
        config['site_dir'] = tempdir
        build(config, live_server=True, clean_site_dir=True)
        return config

    # Perform the initial build
    config = builder()

    host, port = config['dev_addr'].split(':', 1)

    try:
        if livereload:
            _livereload(host, port, config, builder, tempdir)
        else:
            _static_server(host, port, tempdir)
    finally:
        shutil.rmtree(tempdir)
<|endoftext|>


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

import argparse
import logging
import re
import os
import socket
import subprocess

from six.moves.urllib_parse import urlparse

import yaml
import pystache
import requests

def parse_args():
    sections = ['all', 'patroni', 'patronictl', 'certificate', 'wal-e', 'crontab', 'ldap']
    argp = argparse.ArgumentParser(description='Configures Spilo',
                                   epilog="Choose from the following sections:\n\t{}".format('\n\t'.join(sections)),
                                   formatter_class=argparse.RawDescriptionHelpFormatter)

    argp.add_argument('sections', metavar='sections', type=str, nargs='+', choices=sections,
                      help='Which section to (re)configure')
    argp.add_argument('-l', '--loglevel', type=str, help='Explicitly set loglevel')
    argp.add_argument('-f', '--force', help='Overwrite files if they exist', default=False, action='store_true')

    args = vars(argp.parse_args())

    if 'all' in args['sections']:
        args['sections'] = sections
        args['sections'].remove('all')
    args['sections'] = set(args['sections'])

    return args

def write_certificates(environment, overwrite):
    """Write SSL certificate to files

    If certificates are specified, they are written, otherwise
    dummy certificates are generated and written"""

    ssl_keys = ['SSL_CERTIFICATE', 'SSL_PRIVATE_KEY']
    if set(ssl_keys) <= set(environment):
        for k in ssl_keys:
            write_file(environment[k], environment[k + '_FILE'], overwrite)
    else:
        if os.path.exists(environment['SSL_PRIVATE_KEY_FILE']) and not overwrite:
            logging.warning('Private key already exists, not overwriting. (Use option --force if necessary)')
            return
        openssl_cmd = [
            '/usr/bin/openssl',
            'req',
            '-nodes',
            '-new',
            '-x509',
            '-subj',
            '/CN=spilo.dummy.org',
            '-keyout',
            environment['SSL_PRIVATE_KEY_FILE'],
            '-out',
            environment['SSL_CERTIFICATE_FILE'],
        ]
        logging.info('Generating ssl certificate')
        p = subprocess.Popen(openssl_cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        output, _ = p.communicate()
        logging.debug(output)

    uid = os.stat(environment['PGHOME']).st_uid
    os.chmod(environment['SSL_PRIVATE_KEY_FILE'], 0o600)
    os.chown(environment['SSL_PRIVATE_KEY_FILE'], uid, -1)

def deep_update(a, b):
    """Updates data structures

    Dicts are merged, recursively
    List b is appended to a (except duplicates)
    For anything else, the value of a is returned"""

    if type(a) is dict and type(b) is dict:
        for key in b:
            if key in a:
                a[key] = deep_update(a[key], b[key])
            else:
                a[key] = b[key]
        return a
    if type(a) is list and type(b) is list:
        return a + [i for i in b if i not in a]

    return a if a is not None else b

TEMPLATE = \
    '''
ttl: &ttl 30
loop_wait: &loop_wait 10
scope: &scope '{{SCOPE}}'
restapi:
  listen: 0.0.0.0:{{APIPORT}}
  connect_address: {{instance_data.ip}}:{{APIPORT}}
postgresql:
  name: {{instance_data.id}}
  scope: *scope
  listen: 0.0.0.0:{{PGPORT}}
  connect_address: {{instance_data.ip}}:{{PGPORT}}
  data_dir: {{PGDATA}}
  initdb:
    - encoding: UTF8
    - locale: en_US.UTF-8
  pg_hba:
    - hostssl all all 0.0.0.0/0 md5
    - host    all all 0.0.0.0/0 md5
  superuser:
    username: postgres
    password: {{PGPASSWORD_SUPERUSER}}
  admin:
    username: admin
    password: {{PGPASSWORD_ADMIN}}
  replication:
    username: standby
    password: {{PGPASSWORD_STANDBY}}
    network: 0.0.0.0/0
  callbacks:
    on_start: patroni_aws
    on_stop: patroni_aws
    on_restart: patroni_aws
    on_role_change: patroni_aws
  create_replica_method:
    - wal_e
    - basebackup
  wal_e:
    command: patroni_wale_restore
    envdir: {{WALE_ENV_DIR}}
    threshold_megabytes: {{WALE_BACKUP_THRESHOLD_MEGABYTES}}
    threshold_backup_size_percentage: {{WALE_BACKUP_THRESHOLD_PERCENTAGE}}
    use_iam: 1
    retries: 2
  parameters:
    archive_mode: 'on'
    wal_level: hot_standby
    archive_command: envdir "{{WALE_ENV_DIR}}" wal-e --aws-instance-profile wal-push "%p" -p 1
    max_wal_senders: 5
    wal_keep_segments: 8
    archive_timeout: 1800s
    max_connections: {{postgresql.parameters.max_connections}}
    max_replication_slots: 5
    hot_standby: 'on'
    tcp_keepalives_idle: 900
    tcp_keepalives_interval: 100
    ssl: 'on'
    ssl_cert_file: {{SSL_CERTIFICATE_FILE}}
    ssl_key_file: {{SSL_PRIVATE_KEY_FILE}}
    shared_buffers: {{postgresql.parameters.shared_buffers}}
    wal_log_hints: 'on'
    log_line_prefix: '%t [%p]: [%l-1] %c %x %d %u %a %h '
    log_checkpoints: 'on'
    log_lock_waits: 'on'
    log_min_duration_statement: 500
    log_autovacuum_min_duration: 0
    log_connections: 'on'
    log_disconnections: 'on'
    log_statement: 'ddl'
    log_temp_files: 0
  recovery_conf:
    restore_command: envdir "{{WALE_ENV_DIR}}" wal-e --aws-instance-profile wal-fetch "%f" "%p" -p 1
'''

def get_instance_meta_data(key):
    try:
        result = requests.get('http://instance-data/latest/meta-data/{}'.format(key), timeout=2)
        if result.status_code != 200:
            raise Exception('Received status code {} ({})'.format(result.status_code, result.text))
        return result.text
    except Exception as e:
        logging.debug('Could not inspect instance metadata for key {}, error: {}'.format(key, e))

def get_placeholders():
    placeholders = dict(os.environ)

    placeholders.setdefault('PGHOME', os.path.expanduser('~'))
    placeholders.setdefault('APIPORT', '8008')
    placeholders.setdefault('BACKUP_SCHEDULE', '00 01 * * *')
    placeholders.setdefault('CRONTAB', '[]')
    placeholders.setdefault('PGDATA', os.path.join(placeholders['PGHOME'], 'pgdata'))
    placeholders.setdefault('PGPASSWORD_ADMIN', 'standby')
    placeholders.setdefault('PGPASSWORD_STANDBY', 'standby')
    placeholders.setdefault('PGPASSWORD_SUPERUSER', 'zalando')
    placeholders.setdefault('PGPORT', '5432')
    placeholders.setdefault('SCOPE', 'dummy')
    placeholders.setdefault('SSL_CERTIFICATE_FILE', os.path.join(placeholders['PGHOME'], 'server.crt'))
    placeholders.setdefault('SSL_PRIVATE_KEY_FILE', os.path.join(placeholders['PGHOME'], 'server.key'))
    placeholders.setdefault('WALE_BACKUP_THRESHOLD_MEGABYTES', 1024)
    placeholders.setdefault('WALE_BACKUP_THRESHOLD_PERCENTAGE', 30)
    placeholders.setdefault('WALE_ENV_DIR', os.path.join(placeholders['PGHOME'], 'etc', 'wal-e.d', 'env'))
    placeholders.setdefault('WAL_S3_BUCKET', 'spilo-example-com')

    placeholders.setdefault('postgresql', {})
    placeholders['postgresql'].setdefault('parameters', {})

    os_memory_mb = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / 1024 / 1024

    # # We take 1/4 of the memory, expressed in full MB's
    placeholders['postgresql']['parameters']['shared_buffers'] = '{}MB'.format(int(os_memory_mb/4))
    # # 1 connection per 30 MB, at least 100, at most 1000
    placeholders['postgresql']['parameters']['max_connections'] = min(max(100, int(os_memory_mb/30)), 1000)

    placeholders['instance_data'] = dict()
    placeholders['instance_data']['ip'] = get_instance_meta_data('local-ipv4') \
        or socket.gethostbyname(socket.gethostname())
    placeholders['instance_data']['id'] = re.sub(r'\W+', '_', get_instance_meta_data('instance-id')
                                                 or socket.gethostname())

    return placeholders

def write_file(config, filename, overwrite):
    if not overwrite and os.path.exists(filename):
        logging.warning('File {} already exists, not overwriting. (Use option --force if necessary)'.format(filename))
    with open(filename, 'w') as f:
        logging.info('Writing to file {}'.format(filename))
        f.write(config)

def pystache_render(*args, **kwargs):
    render = pystache.Renderer(missing_tags='strict')
    return render.render(*args, **kwargs)

def get_dcs_config(config, placeholders):
    defaults = \
        yaml.load('''\
zookeeper:
  scope: {scope}
  session_timeout: {ttl}
  reconnect_timeout: {loop_wait}
etcd:
  scope: {scope}
  ttl: {ttl}'''.format(**config))

    config = {}

    if 'ZOOKEEPER_HOSTS' in placeholders:
        config = {'zookeeper': defaults['zookeeper']}
        config['zookeeper']['hosts'] = yaml.load(placeholders['ZOOKEEPER_HOSTS'])
    elif 'EXHIBITOR_HOSTS' in placeholders and 'EXHIBITOR_PORT' in placeholders:
        config = {'zookeeper': defaults['zookeeper']}
        config['zookeeper']['exhibitor'] = {'poll_interval': '300', 'port': placeholders['EXHIBITOR_PORT'],
                                            'hosts': yaml.load(placeholders['EXHIBITOR_HOSTS'])}
    elif 'ETCD_HOST' in placeholders:
        config = {'etcd': defaults['etcd']}
        config['etcd']['host'] = placeholders['ETCD_HOST']
    elif 'ETCD_DISCOVERY_DOMAIN' in placeholders:
        config = {'etcd': defaults['etcd']}
        config['etcd']['discovery_srv'] = placeholders['ETCD_DISCOVERY_DOMAIN']
    else:
        pass  # Configuration can also be specified using PATRONI_CONFIGURATION

    return config

def write_wale_command_environment(placeholders, overwrite):
    az = get_instance_meta_data('placement/availability-zone') or 'dummy-region'

    if not os.path.exists(placeholders['WALE_ENV_DIR']):
        os.makedirs(placeholders['WALE_ENV_DIR'])

    write_file('s3://{WAL_S3_BUCKET}/spilo/{SCOPE}/wal/'.format(**placeholders),
               os.path.join(placeholders['WALE_ENV_DIR'], 'WALE_S3_PREFIX'), overwrite)
    write_file('https+path://s3-{}.amazonaws.com:443'.format(az[:-1]),
               os.path.join(placeholders['WALE_ENV_DIR'], 'WALE_S3_ENDPOINT'), overwrite)

def write_crontab(placeholders, path, overwrite):

    if not overwrite:
        with open(os.devnull, 'w') as devnull:
            cron_exit = subprocess.call(['sudo', '-u', 'postgres', 'crontab', '-l'], stdout=devnull, stderr=devnull)
            if cron_exit == 0:
                logging.warning('Cron is already configured. (Use option --force to overwrite cron)')
                return

    lines = ['PATH={}'.format(path)]
    lines += ['{BACKUP_SCHEDULE} /postgres_backup.sh "{WALE_ENV_DIR}" "{PGDATA}"'.format(**placeholders)]
    lines += yaml.load(placeholders['CRONTAB'])
    lines += ['']  # EOF requires empty line for cron

    c = subprocess.Popen(['sudo', '-u', 'postgres', 'crontab'], stdin=subprocess.PIPE)
    c.communicate(input='\n'.join(lines).encode())

def write_ldap_configuration(placeholders, overwrite):
    ldap_url = placeholders.get('LDAP_URL')
    if ldap_url is None:
        logging.info("No LDAP_URL was specified, skipping LDAP configuration")
        return

    r = urlparse(ldap_url)
    if not r.scheme:
        logging.error('LDAP_URL should contain a scheme')
        logging.info(r)
        return

    host, port = r.hostname, r.port
    if not port:
        port = 636 if r.scheme == 'ldaps' else 389

    stunnel_config = """\
foreground = yes
options = NO_SSLv2

[ldaps]
connect = {0}:{1}
client = yes
accept = 389
verify = 3
CAfile = /etc/stunnel/chain.pem
""".format(host, port)
    write_file(stunnel_config, '/etc/stunnel/ldap.conf', overwrite)

    supervisord_config = """\
[program:ldaptunnel]
autostart=true
priority=500
directory=/
command=env -i /usr/bin/stunnel4 /etc/stunnel/ldap.conf
stdout_logfile=/dev/stdout
stdout_logfile_maxbytes=0
redirect_stderr=true
"""
    write_file(supervisord_config, '/etc/supervisor/conf.d/ldaptunnel.conf', overwrite)

def main():
    debug = os.environ.get('DEBUG', '') in ['1', 'true', 'on', 'ON']
    args = parse_args()

    logging.basicConfig(format='%(asctime)s - bootstrapping - %(levelname)s - %(message)s', level=('DEBUG'
                         if debug else (args.get('loglevel') or 'INFO').upper()))

    placeholders = get_placeholders()

    config = yaml.load(pystache_render(TEMPLATE, placeholders))
    config.update(get_dcs_config(config, placeholders))

    user_config = yaml.load(os.environ.get('PATRONI_CONFIGURATION', '')) or {}
    if not isinstance(user_config, dict):
        raise ValueError('PATRONI_CONFIGURATION should contain a dict, yet it is a {}'.format(type(user_config)))

    config = deep_update(user_config, config)

    # Ensure replication is available
    if not any(['replication' in i for i in config['postgresql']['pg_hba']]):
        rep_hba = 'hostssl replication {} 0.0.0.0/0 md5'.format(config['postgresql']['replication']['username'])
        config['postgresql']['pg_hba'].insert(0, rep_hba)

    for section in args['sections']:
        logging.info('Configuring {}'.format(section))
        if section == 'patroni':
            patroni_configfile = os.path.join(placeholders['PGHOME'], 'postgres.yml')
            write_file(yaml.dump(config, default_flow_style=False, width=120), patroni_configfile, args['force'])
        elif section == 'patronictl':
            patronictl_config = {k: v for k, v in config.items() if k in ['zookeeper', 'etcd', 'consul']}
            patronictl_configfile = os.path.join(placeholders['PGHOME'], '.config', 'patroni', 'patronictl.yaml')
            if not os.path.exists(os.path.dirname(patronictl_configfile)):
                os.makedirs(os.path.dirname(patronictl_configfile))
            write_file(yaml.dump(patronictl_config), patronictl_configfile, args['force'])
        elif section == 'wal-e':
            write_wale_command_environment(placeholders, args['force'])
        elif section == 'certificate':
            write_certificates(placeholders, args['force'])
        elif section == 'crontab':
            write_crontab(placeholders, os.environ.get('PATH'), args['force'])
        elif section == 'ldap':
            write_ldap_configuration(placeholders, args['force'])
        else:
            raise Exception('Unknown section: {}'.format(section))

if __name__ == '__main__':
    main()
<commit_msg>Allow restore without a master<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import re
import os
import socket
import subprocess

from six.moves.urllib_parse import urlparse

import yaml
import pystache
import requests

def parse_args():
    sections = ['all', 'patroni', 'patronictl', 'certificate', 'wal-e', 'crontab', 'ldap']
    argp = argparse.ArgumentParser(description='Configures Spilo',
                                   epilog="Choose from the following sections:\n\t{}".format('\n\t'.join(sections)),
                                   formatter_class=argparse.RawDescriptionHelpFormatter)

    argp.add_argument('sections', metavar='sections', type=str, nargs='+', choices=sections,
                      help='Which section to (re)configure')
    argp.add_argument('-l', '--loglevel', type=str, help='Explicitly set loglevel')
    argp.add_argument('-f', '--force', help='Overwrite files if they exist', default=False, action='store_true')

    args = vars(argp.parse_args())

    if 'all' in args['sections']:
        args['sections'] = sections
        args['sections'].remove('all')
    args['sections'] = set(args['sections'])

    return args

def write_certificates(environment, overwrite):
    """Write SSL certificate to files

    If certificates are specified, they are written, otherwise
    dummy certificates are generated and written"""

    ssl_keys = ['SSL_CERTIFICATE', 'SSL_PRIVATE_KEY']
    if set(ssl_keys) <= set(environment):
        for k in ssl_keys:
            write_file(environment[k], environment[k + '_FILE'], overwrite)
    else:
        if os.path.exists(environment['SSL_PRIVATE_KEY_FILE']) and not overwrite:
            logging.warning('Private key already exists, not overwriting. (Use option --force if necessary)')
            return
        openssl_cmd = [
            '/usr/bin/openssl',
            'req',
            '-nodes',
            '-new',
            '-x509',
            '-subj',
            '/CN=spilo.dummy.org',
            '-keyout',
            environment['SSL_PRIVATE_KEY_FILE'],
            '-out',
            environment['SSL_CERTIFICATE_FILE'],
        ]
        logging.info('Generating ssl certificate')
        p = subprocess.Popen(openssl_cmd, shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
        output, _ = p.communicate()
        logging.debug(output)

    uid = os.stat(environment['PGHOME']).st_uid
    os.chmod(environment['SSL_PRIVATE_KEY_FILE'], 0o600)
    os.chown(environment['SSL_PRIVATE_KEY_FILE'], uid, -1)

def deep_update(a, b):
    """Updates data structures

    Dicts are merged, recursively
    List b is appended to a (except duplicates)
    For anything else, the value of a is returned"""

    if type(a) is dict and type(b) is dict:
        for key in b:
            if key in a:
                a[key] = deep_update(a[key], b[key])
            else:
                a[key] = b[key]
        return a
    if type(a) is list and type(b) is list:
        return a + [i for i in b if i not in a]

    return a if a is not None else b

TEMPLATE = \
    '''
ttl: &ttl 30
loop_wait: &loop_wait 10
scope: &scope '{{SCOPE}}'
restapi:
  listen: 0.0.0.0:{{APIPORT}}
  connect_address: {{instance_data.ip}}:{{APIPORT}}
postgresql:
  name: {{instance_data.id}}
  scope: *scope
  listen: 0.0.0.0:{{PGPORT}}
  connect_address: {{instance_data.ip}}:{{PGPORT}}
  data_dir: {{PGDATA}}
  initdb:
    - encoding: UTF8
    - locale: en_US.UTF-8
  pg_hba:
    - hostssl all all 0.0.0.0/0 md5
    - host    all all 0.0.0.0/0 md5
  superuser:
    username: postgres
    password: {{PGPASSWORD_SUPERUSER}}
  admin:
    username: admin
    password: {{PGPASSWORD_ADMIN}}
  replication:
    username: standby
    password: {{PGPASSWORD_STANDBY}}
    network: 0.0.0.0/0
  callbacks:
    on_start: patroni_aws
    on_stop: patroni_aws
    on_restart: patroni_aws
    on_role_change: patroni_aws
  create_replica_method:
    - wal_e
    - basebackup
  wal_e:
    command: patroni_wale_restore
    envdir: {{WALE_ENV_DIR}}
    threshold_megabytes: {{WALE_BACKUP_THRESHOLD_MEGABYTES}}
    threshold_backup_size_percentage: {{WALE_BACKUP_THRESHOLD_PERCENTAGE}}
    use_iam: 1
    retries: 2
    no_master: 1
  parameters:
    archive_mode: 'on'
    wal_level: hot_standby
    archive_command: envdir "{{WALE_ENV_DIR}}" wal-e --aws-instance-profile wal-push "%p" -p 1
    max_wal_senders: 5
    wal_keep_segments: 8
    archive_timeout: 1800s
    max_connections: {{postgresql.parameters.max_connections}}
    max_replication_slots: 5
    hot_standby: 'on'
    tcp_keepalives_idle: 900
    tcp_keepalives_interval: 100
    ssl: 'on'
    ssl_cert_file: {{SSL_CERTIFICATE_FILE}}
    ssl_key_file: {{SSL_PRIVATE_KEY_FILE}}
    shared_buffers: {{postgresql.parameters.shared_buffers}}
    wal_log_hints: 'on'
    log_line_prefix: '%t [%p]: [%l-1] %c %x %d %u %a %h '
    log_checkpoints: 'on'
    log_lock_waits: 'on'
    log_min_duration_statement: 500
    log_autovacuum_min_duration: 0
    log_connections: 'on'
    log_disconnections: 'on'
    log_statement: 'ddl'
    log_temp_files: 0
  recovery_conf:
    restore_command: envdir "{{WALE_ENV_DIR}}" wal-e --aws-instance-profile wal-fetch "%f" "%p" -p 1
'''

def get_instance_meta_data(key):
    try:
        result = requests.get('http://instance-data/latest/meta-data/{}'.format(key), timeout=2)
        if result.status_code != 200:
            raise Exception('Received status code {} ({})'.format(result.status_code, result.text))
        return result.text
    except Exception as e:
        logging.debug('Could not inspect instance metadata for key {}, error: {}'.format(key, e))

def get_placeholders():
    placeholders = dict(os.environ)

    placeholders.setdefault('PGHOME', os.path.expanduser('~'))
    placeholders.setdefault('APIPORT', '8008')
    placeholders.setdefault('BACKUP_SCHEDULE', '00 01 * * *')
    placeholders.setdefault('CRONTAB', '[]')
    placeholders.setdefault('PGDATA', os.path.join(placeholders['PGHOME'], 'pgdata'))
    placeholders.setdefault('PGPASSWORD_ADMIN', 'standby')
    placeholders.setdefault('PGPASSWORD_STANDBY', 'standby')
    placeholders.setdefault('PGPASSWORD_SUPERUSER', 'zalando')
    placeholders.setdefault('PGPORT', '5432')
    placeholders.setdefault('SCOPE', 'dummy')
    placeholders.setdefault('SSL_CERTIFICATE_FILE', os.path.join(placeholders['PGHOME'], 'server.crt'))
    placeholders.setdefault('SSL_PRIVATE_KEY_FILE', os.path.join(placeholders['PGHOME'], 'server.key'))
    placeholders.setdefault('WALE_BACKUP_THRESHOLD_MEGABYTES', 1024)
    placeholders.setdefault('WALE_BACKUP_THRESHOLD_PERCENTAGE', 30)
    placeholders.setdefault('WALE_ENV_DIR', os.path.join(placeholders['PGHOME'], 'etc', 'wal-e.d', 'env'))
    placeholders.setdefault('WAL_S3_BUCKET', 'spilo-example-com')

    placeholders.setdefault('postgresql', {})
    placeholders['postgresql'].setdefault('parameters', {})

    os_memory_mb = os.sysconf('SC_PAGE_SIZE') * os.sysconf('SC_PHYS_PAGES') / 1024 / 1024

    # # We take 1/4 of the memory, expressed in full MB's
    placeholders['postgresql']['parameters']['shared_buffers'] = '{}MB'.format(int(os_memory_mb/4))
    # # 1 connection per 30 MB, at least 100, at most 1000
    placeholders['postgresql']['parameters']['max_connections'] = min(max(100, int(os_memory_mb/30)), 1000)

    placeholders['instance_data'] = dict()
    placeholders['instance_data']['ip'] = get_instance_meta_data('local-ipv4') \
        or socket.gethostbyname(socket.gethostname())
    placeholders['instance_data']['id'] = re.sub(r'\W+', '_', get_instance_meta_data('instance-id')
                                                 or socket.gethostname())

    return placeholders

def write_file(config, filename, overwrite):
    if not overwrite and os.path.exists(filename):
        logging.warning('File {} already exists, not overwriting. (Use option --force if necessary)'.format(filename))
    with open(filename, 'w') as f:
        logging.info('Writing to file {}'.format(filename))
        f.write(config)

def pystache_render(*args, **kwargs):
    render = pystache.Renderer(missing_tags='strict')
    return render.render(*args, **kwargs)

def get_dcs_config(config, placeholders):
    defaults = \
        yaml.load('''\
zookeeper:
  scope: {scope}
  session_timeout: {ttl}
  reconnect_timeout: {loop_wait}
etcd:
  scope: {scope}
  ttl: {ttl}'''.format(**config))

    config = {}

    if 'ZOOKEEPER_HOSTS' in placeholders:
        config = {'zookeeper': defaults['zookeeper']}
        config['zookeeper']['hosts'] = yaml.load(placeholders['ZOOKEEPER_HOSTS'])
    elif 'EXHIBITOR_HOSTS' in placeholders and 'EXHIBITOR_PORT' in placeholders:
        config = {'zookeeper': defaults['zookeeper']}
        config['zookeeper']['exhibitor'] = {'poll_interval': '300', 'port': placeholders['EXHIBITOR_PORT'],
                                            'hosts': yaml.load(placeholders['EXHIBITOR_HOSTS'])}
    elif 'ETCD_HOST' in placeholders:
        config = {'etcd': defaults['etcd']}
        config['etcd']['host'] = placeholders['ETCD_HOST']
    elif 'ETCD_DISCOVERY_DOMAIN' in placeholders:
        config = {'etcd': defaults['etcd']}
        config['etcd']['discovery_srv'] = placeholders['ETCD_DISCOVERY_DOMAIN']
    else:
        pass  # Configuration can also be specified using PATRONI_CONFIGURATION

    return config

def write_wale_command_environment(placeholders, overwrite):
    az = get_instance_meta_data('placement/availability-zone') or 'dummy-region'

    if not os.path.exists(placeholders['WALE_ENV_DIR']):
        os.makedirs(placeholders['WALE_ENV_DIR'])

    write_file('s3://{WAL_S3_BUCKET}/spilo/{SCOPE}/wal/'.format(**placeholders),
               os.path.join(placeholders['WALE_ENV_DIR'], 'WALE_S3_PREFIX'), overwrite)
    write_file('https+path://s3-{}.amazonaws.com:443'.format(az[:-1]),
               os.path.join(placeholders['WALE_ENV_DIR'], 'WALE_S3_ENDPOINT'), overwrite)

def write_crontab(placeholders, path, overwrite):

    if not overwrite:
        with open(os.devnull, 'w') as devnull:
            cron_exit = subprocess.call(['sudo', '-u', 'postgres', 'crontab', '-l'], stdout=devnull, stderr=devnull)
            if cron_exit == 0:
                logging.warning('Cron is already configured. (Use option --force to overwrite cron)')
                return

    lines = ['PATH={}'.format(path)]
    lines += ['{BACKUP_SCHEDULE} /postgres_backup.sh "{WALE_ENV_DIR}" "{PGDATA}"'.format(**placeholders)]
    lines += yaml.load(placeholders['CRONTAB'])
    lines += ['']  # EOF requires empty line for cron

    c = subprocess.Popen(['sudo', '-u', 'postgres', 'crontab'], stdin=subprocess.PIPE)
    c.communicate(input='\n'.join(lines).encode())

def write_ldap_configuration(placeholders, overwrite):
    ldap_url = placeholders.get('LDAP_URL')
    if ldap_url is None:
        logging.info("No LDAP_URL was specified, skipping LDAP configuration")
        return

    r = urlparse(ldap_url)
    if not r.scheme:
        logging.error('LDAP_URL should contain a scheme')
        logging.info(r)
        return

    host, port = r.hostname, r.port
    if not port:
        port = 636 if r.scheme == 'ldaps' else 389

    stunnel_config = """\
foreground = yes
options = NO_SSLv2

[ldaps]
connect = {0}:{1}
client = yes
accept = 389
verify = 3
CAfile = /etc/stunnel/chain.pem
""".format(host, port)
    write_file(stunnel_config, '/etc/stunnel/ldap.conf', overwrite)

    supervisord_config = """\
[program:ldaptunnel]
autostart=true
priority=500
directory=/
command=env -i /usr/bin/stunnel4 /etc/stunnel/ldap.conf
stdout_logfile=/dev/stdout
stdout_logfile_maxbytes=0
redirect_stderr=true
"""
    write_file(supervisord_config, '/etc/supervisor/conf.d/ldaptunnel.conf', overwrite)

def main():
    debug = os.environ.get('DEBUG', '') in ['1', 'true', 'on', 'ON']
    args = parse_args()

    logging.basicConfig(format='%(asctime)s - bootstrapping - %(levelname)s - %(message)s', level=('DEBUG'
                         if debug else (args.get('loglevel') or 'INFO').upper()))

    placeholders = get_placeholders()

    config = yaml.load(pystache_render(TEMPLATE, placeholders))
    config.update(get_dcs_config(config, placeholders))

    user_config = yaml.load(os.environ.get('PATRONI_CONFIGURATION', '')) or {}
    if not isinstance(user_config, dict):
        raise ValueError('PATRONI_CONFIGURATION should contain a dict, yet it is a {}'.format(type(user_config)))

    config = deep_update(user_config, config)

    # Ensure replication is available
    if not any(['replication' in i for i in config['postgresql']['pg_hba']]):
        rep_hba = 'hostssl replication {} 0.0.0.0/0 md5'.format(config['postgresql']['replication']['username'])
        config['postgresql']['pg_hba'].insert(0, rep_hba)

    for section in args['sections']:
        logging.info('Configuring {}'.format(section))
        if section == 'patroni':
            patroni_configfile = os.path.join(placeholders['PGHOME'], 'postgres.yml')
            write_file(yaml.dump(config, default_flow_style=False, width=120), patroni_configfile, args['force'])
        elif section == 'patronictl':
            patronictl_config = {k: v for k, v in config.items() if k in ['zookeeper', 'etcd', 'consul']}
            patronictl_configfile = os.path.join(placeholders['PGHOME'], '.config', 'patroni', 'patronictl.yaml')
            if not os.path.exists(os.path.dirname(patronictl_configfile)):
                os.makedirs(os.path.dirname(patronictl_configfile))
            write_file(yaml.dump(patronictl_config), patronictl_configfile, args['force'])
        elif section == 'wal-e':
            write_wale_command_environment(placeholders, args['force'])
        elif section == 'certificate':
            write_certificates(placeholders, args['force'])
        elif section == 'crontab':
            write_crontab(placeholders, os.environ.get('PATH'), args['force'])
        elif section == 'ldap':
            write_ldap_configuration(placeholders, args['force'])
        else:
            raise Exception('Unknown section: {}'.format(section))

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals
# Erik Stein , 2015-2017

import html
import re

from django.utils import six
from django.utils.encoding import force_text
from django.utils.functional import keep_lazy, keep_lazy_text
from django.utils.safestring import SafeText
from django.utils.html import mark_safe
from django.utils.text import slugify
from django.utils.translation import ugettext_lazy

# from bs4 import BeautifulStoneSoup
import translitcodec  # provides 'translit/long', used by codecs.encode()
import codecs

def downgrade(value):
    """
    Downgrade unicode to ascii, transliterating accented characters.
    """
    value = force_text(value)
    return codecs.encode(value, 'translit/long')
# downgrade = allow_lazy(downgrade, six.text_type, SafeText)

def slugify_long(value):
    return slugify(downgrade(value))
# slugify_long = allow_lazy(slugify_long, six.text_type, SafeText)

# Spreading umlauts is included in the translit/long codec.
slugify_german = slugify_long

# Does not work anymore with bs4
# def html_entities_to_unicode(html):
#     text = smart_text(BeautifulStoneSoup(html, convertEntities=BeautifulStoneSoup.ALL_ENTITIES))
#     return text

# Works only with Python >= 3.4
def html_entities_to_unicode(html_str):
    return html.unescape(html_str)
# html_entities_to_unicode = allow_lazy(html_entities_to_unicode, six.text_type, SafeText)

# Translators: This string is used as a separator between list elements
DEFAULT_SEPARATOR = ugettext_lazy(", ")

@keep_lazy_text
def get_text_joined(list_, separator=DEFAULT_SEPARATOR, last_word=ugettext_lazy(' and ')):
    list_ = list(list_)
    if len(list_) == 0:
        return ''
    if len(list_) == 1:
        return force_text(list_[0])
    return '%s%s%s' % (
        separator.join(force_text(i) for i in list_[:-1]),
        force_text(last_word), force_text(list_[-1]))

def slimdown(text):
    """
    Converts simplified markdown (**, *, _) to ,  und  tags.
    """
    b_pattern = re.compile(r"(\*\*)(.*?)\1")
    i_pattern = re.compile(r"(\*)(.*?)\1")
    u_pattern = re.compile(r"(__)(.*?)\1")

    text, n = re.subn(b_pattern, "\\2", text)
    text, n = re.subn(i_pattern, "\\2", text)
    text, n = re.subn(u_pattern, "\\2", text)
    return mark_safe(text)

<commit_msg>htmlentities_to_unicode Python 2 compatibility.<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals
# Erik Stein , 2015-2017

from django.utils import six
if six.PY3:
    import html

import re

from django.utils.encoding import force_text, smart_text
from django.utils.functional import keep_lazy, keep_lazy_text
from django.utils.safestring import SafeText
from django.utils.html import mark_safe
from django.utils.text import slugify
from django.utils.translation import ugettext_lazy

# from bs4 import BeautifulStoneSoup
import translitcodec  # provides 'translit/long', used by codecs.encode()
import codecs

def downgrade(value):
    """
    Downgrade unicode to ascii, transliterating accented characters.
    """
    value = force_text(value)
    return codecs.encode(value, 'translit/long')
# downgrade = allow_lazy(downgrade, six.text_type, SafeText)

def slugify_long(value):
    return slugify(downgrade(value))
# slugify_long = allow_lazy(slugify_long, six.text_type, SafeText)

# Spreading umlauts is included in the translit/long codec.
slugify_german = slugify_long

if six.PY2:
    import bs4

    def html_entities_to_unicode(html):
        # An incoming HTML or XML entity is always converted into the corresponding Unicode character in bs4
        return smart_text(bs4.BeautifulStoneSoup(html))

else:
    # Works only with Python >= 3.4
    def html_entities_to_unicode(html_str):
        return html.unescape(html_str)
    # html_entities_to_unicode = allow_lazy(html_entities_to_unicode, six.text_type, SafeText)

# Translators: This string is used as a separator between list elements
DEFAULT_SEPARATOR = ugettext_lazy(", ")

@keep_lazy_text
def get_text_joined(list_, separator=DEFAULT_SEPARATOR, last_word=ugettext_lazy(' and ')):
    list_ = list(list_)
    if len(list_) == 0:
        return ''
    if len(list_) == 1:
        return force_text(list_[0])
    return '%s%s%s' % (
        separator.join(force_text(i) for i in list_[:-1]),
        force_text(last_word), force_text(list_[-1]))

def slimdown(text):
    """
    Converts simplified markdown (**, *, _) to ,  und  tags.
    """
    b_pattern = re.compile(r"(\*\*)(.*?)\1")
    i_pattern = re.compile(r"(\*)(.*?)\1")
    u_pattern = re.compile(r"(__)(.*?)\1")

    text, n = re.subn(b_pattern, "\\2", text)
    text, n = re.subn(i_pattern, "\\2", text)
    text, n = re.subn(u_pattern, "\\2", text)
    return mark_safe(text)

<|endoftext|>


<commit_before>import pytest
import rml.lattice
import rml.element
import mock

DUMMY_NAME = 'dummy'

@pytest.fixture
def simple_element():
    e_length = 1.5
    e = rml.element.Element('dummy_element', 'Quad', mock.MagicMock())
    e.add_to_family('Q1B')
    return e

@pytest.fixture
def simple_element_and_lattice(simple_element):
    l = rml.lattice.Lattice(DUMMY_NAME)
    l.add_element(simple_element)
    return simple_element, l

def test_create_lattice():
    l = rml.lattice.Lattice('DUMMY_NAME')
    assert(len(l)) == 0
    assert l.name == 'DUMMY_NAME'

def test_non_negative_lattice():
    l = rml.lattice.Lattice(DUMMY_NAME)
    assert(len(l)) >= 0

def test_lattice_with_n_elements(simple_element_and_lattice):
    elem, lattice = simple_element_and_lattice

    # Getting elements
    lattice.add_element(elem)
    assert lattice[0] == elem
    assert lattice.get_elements() == [elem, elem]

def test_lattice_get_element_with_family(simple_element_and_lattice):
    element, lattice = simple_element_and_lattice
    element.add_to_family('fam')
    assert lattice.get_elements('fam') == [element]
    assert lattice.get_elements('nofam') == []

def test_get_all_families(simple_element_and_lattice):
    element, lattice = simple_element_and_lattice
    families = lattice.get_all_families()
    assert len(families) > 0
<commit_msg>Tests for setting the pv value for a whole family<commit_after>import pytest
import rml.lattice
import rml.element
import rml.device
import cs_dummy
import mock
from rml.units import UcPoly

DUMMY_NAME = 'dummy'

@pytest.fixture
def simple_element(identity=1):
    cs = cs_dummy.CsDummy()
    uc = UcPoly([0, 1])

    # Create devices and attach them to the element
    element = rml.element.Element(identity, 'BPM', mock.MagicMock())
    rb_pv = 'SR22C-DI-EBPM-04:SA:X'
    sp_pv = 'SR22C-DI-EBPM-04:SA:Y'
    device1 = rml.device.Device(rb_pv, sp_pv, cs)
    device2 = rml.device.Device(sp_pv, rb_pv, cs)
    element.add_to_family('BPM')

    element.add_device('x', device1, uc)
    element.add_device('y', device2, uc)

    return element

@pytest.fixture
def simple_element_and_lattice(simple_element):
    l = rml.lattice.Lattice(DUMMY_NAME)
    l.add_element(simple_element)
    return simple_element, l

def test_create_lattice():
    l = rml.lattice.Lattice('DUMMY_NAME')
    assert(len(l)) == 0
    assert l.name == 'DUMMY_NAME'

def test_non_negative_lattice():
    l = rml.lattice.Lattice(DUMMY_NAME)
    assert(len(l)) >= 0

def test_lattice_with_n_elements(simple_element_and_lattice):
    elem, lattice = simple_element_and_lattice

    # Getting elements
    lattice.add_element(elem)
    assert lattice[0] == elem
    assert lattice.get_elements() == [elem, elem]

def test_lattice_get_element_with_family(simple_element_and_lattice):
    element, lattice = simple_element_and_lattice
    element.add_to_family('fam')
    assert lattice.get_elements('fam') == [element]
    assert lattice.get_elements('nofam') == []

def test_get_all_families(simple_element_and_lattice):
    element, lattice = simple_element_and_lattice
    families = lattice.get_all_families()
    assert len(families) > 0

def test_get_family_value(simple_element_and_lattice):
    element, lattice = simple_element_and_lattice
    lattice.set_family_value('BPM', 'x', [2.3])
    family_values = lattice.get_family_value('BPM', 'x')
    assert family_values == list([2.3])
<|endoftext|>


<commit_before>import rml.lattice

def test_create_lattice():
    l = rml.lattice.Lattice()
    assert(len(l)) == 0

<commit_msg>Test lattince length for non-negative values<commit_after>import rml.lattice

def test_create_lattice():
    l = rml.lattice.Lattice()
    assert(len(l)) == 0

def test_non_negative_lattice():
    l = rml.lattice.Lattice()
    assert(len(l)) >= 0
<|endoftext|>


<commit_before>""" Metrics Tests"""
import os
import tempfile
import time
import unittest
from pathlib import Path
from typing import List

import torch
from torch import nn

from avalanche.evaluation import PluginMetric
from avalanche.evaluation.metrics import Accuracy, Loss, ConfusionMatrix, \
    DiskUsage, MAC, CPUUsage, MaxGPU, MaxRAM, Mean, Sum, ElapsedTime, Forgetting
from avalanche.evaluation.metrics import accuracy_metrics, loss_metrics, StreamConfusionMatrix, \
    disk_usage_metrics, MAC_metrics, cpu_usage_metrics, gpu_usage_metrics, ram_usage_metrics, \
    timing_metrics
from tests.unit_tests_utils import common_setups

#################################
#################################
##### GENERAL METRIC TEST #######
#################################
#################################

class GeneralMetricTests(unittest.TestCase):
    def setUp(self) -> None:
        self.batch_size = 3
        self.input_size = 10
        self.n_classes = 3
        self.out = torch.randn(self.batch_size, self.input_size)
        self.y = torch.randint(0, self.n_classes, (self.batch_size,))

    def test_accuracy(self):
        metric = Accuracy()
        self.assertEqual(metric.result(), 0)
        metric.update(self.out, self.y)
        self.assertLessEqual(metric.result(), 1)
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_loss(self):
        metric = Loss()
        self.assertEqual(metric.result(), 0)
        metric.update(torch.tensor(1.), self.batch_size)
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_cm(self):
        metric = ConfusionMatrix()
        cm = metric.result()
        self.assertTrue((cm == 0).all().item())
        metric.update(self.y, self.out)
        cm = metric.result()
        self.assertTrue((cm >= 0).all().item())
        metric.reset()
        cm = metric.result()
        self.assertTrue((cm == 0).all().item())

    def test_ram(self):
        metric = MaxRAM()
        self.assertEqual(metric.result(), 0)
        metric.start_thread()  # start thread
        self.assertGreaterEqual(metric.result(), 0)
        metric.stop_thread() # stop thread
        metric.reset()  # stop thread
        self.assertEqual(metric.result(), 0)

    def test_gpu(self):
        if torch.cuda.is_available():
            metric = MaxGPU(0)
            self.assertEqual(metric.result(), 0)
            metric.start_thread()  # start thread
            self.assertGreaterEqual(metric.result(), 0)
            metric.stop_thread() # stop thread
            metric.reset()  # stop thread
            self.assertEqual(metric.result(), 0)

    def test_cpu(self):
        metric = CPUUsage()
        self.assertEqual(metric.result(), 0)
        metric.update()
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_disk(self):
        metric = DiskUsage()
        self.assertEqual(metric.result(), 0)
        metric.update()
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_timing(self):
        metric = ElapsedTime()
        self.assertEqual(metric.result(), 0)
        metric.update()  # need two update calls
        self.assertEqual(metric.result(), 0)
        metric.update()
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_mac(self):
        model = torch.nn.Linear(self.input_size, 2)
        metric = MAC()
        self.assertEqual(metric.result(), 0)
        metric.update(model, self.out)
        self.assertGreaterEqual(metric.result(), 0)

    def test_mean(self):
        metric = Mean()
        self.assertEqual(metric.result(), 0)
        metric.update(0.1, 1)
        self.assertEqual(metric.result(), 0.1)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_sum(self):
        metric = Sum()
        self.assertEqual(metric.result(), 0)
        metric.update(5)
        self.assertEqual(metric.result(), 5)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_forgetting(self):
        metric = Forgetting()
        f = metric.result()
        self.assertEqual(f, {})
        f = metric.result(k=0)
        self.assertIsNone(f)
        metric.update(0, 1, initial=True)
        f = metric.result(k=0)
        self.assertIsNone(f)
        metric.update(0, 0.4)
        f = metric.result(k=0)
        self.assertEqual(f, 0.6)
        metric.reset()
        self.assertEqual(metric.result(), {})

#################################
#################################
### FINE GRAINED METRIC TEST ####
#################################
#################################

class MACMetricTests(unittest.TestCase):
    def setUp(self):
        common_setups()

    def test_ff_model(self):
        xn, hn, yn = 50, 100, 10

        model = nn.Sequential(
            nn.Linear(xn, hn),  # 5'000 mul
            nn.ReLU(),
            nn.Linear(hn, hn),  # 10'000 mul
            nn.ReLU(),
            nn.Linear(hn, yn)  # 1'000 mul
        )  # 16'000 mul
        dummy = torch.randn(32, xn)
        met = MAC()
        # met.update(model, dummy)
        # self.assertEqual(16000, met.result())

    def test_cnn_model(self):
        xn, hn, yn = 50, 100, 10
        model = nn.Sequential(
            nn.Conv2d(1, 10, 4),  # 353'440 mul
            nn.ReLU(),
            nn.Flatten(),
            nn.Linear(22090, hn),  # 2'209'000 mul
            nn.Linear(hn, yn)  # 1'000 mul
        )  # 2'563'440 mul
        dummy = torch.randn(32, 1, xn, xn)
        met = MAC()
        # print(2563440, met.compute(model, dummy))

class AccuracyMetricTests(unittest.TestCase):
    def setUp(self):
        common_setups()

    def test_standalone_accuracy(self):
        uut = Accuracy()

        # Initial accuracy should be 0
        self.assertEqual(0.0, uut.result())

        truth = torch.as_tensor([0, 5, 2, 1, 0])
        predicted = torch.as_tensor([2, 3, 2, 5, 0])  # correct 2/5 = 40%
        uut.update(truth, predicted)

        self.assertEqual(0.4, uut.result())

        truth = torch.as_tensor([0, 3, 2, 1, 0])
        predicted = torch.as_tensor([2, 3, 2, 5, 0])  # correct 3/5 = 60%
        uut.update(truth, predicted)

        self.assertEqual(0.5, uut.result())

        # After-reset accuracy should be 0
        uut.reset()
        self.assertEqual(0.0, uut.result())

        # Check if handles 0 accuracy
        truth = torch.as_tensor([0, 0, 0, 0])
        predicted = torch.as_tensor([1, 1, 1, 1])  # correct 0/4 = 0%
        uut.update(truth, predicted)

        self.assertEqual(0.0, uut.result())

        # Should throw exception when len(truth) != len(predicted)
        with self.assertRaises(ValueError):
            truth = torch.as_tensor([0, 0, 1, 0])
            predicted = torch.as_tensor([1, 1, 1])
            uut.update(truth, predicted)

        # Check accuracy didn't change after error
        self.assertEqual(0.0, uut.result())

        # Test logits / one-hot support
        uut.reset()

        # Test one-hot (truth)
        truth = torch.as_tensor([[0, 1], [1, 0], [1, 0], [1, 0]])
        predicted = torch.as_tensor([1, 1, 1, 1])  # correct 1/4 = 25%
        uut.update(truth, predicted)
        self.assertEqual(0.25, uut.result())

        # Test logits (predictions)
        truth = torch.as_tensor([1, 1, 0, 0])
        predicted = torch.as_tensor(
            [[0.73, 0.1], [0.22, 0.33],
             [0.99, 0.01], [0.12, 0.11]])  # correct 3/4 = 75%
        uut.update(truth, predicted)
        self.assertEqual(0.5, uut.result())

        # Test one-hot (truth) + logits (predictions)
        truth = torch.as_tensor([[1, 0], [1, 0], [0, 1], [0, 1]])
        predicted = torch.as_tensor(
            [[0.73, 0.1], [0.22, 0.33],
             [0.99, 0.01], [0.12, 0.11]])  # correct 1/4 = 25%
        uut.update(truth, predicted)
        self.assertEqual(5.0/12.0, uut.result())

    def test_accuracy_helper(self):
        metrics = accuracy_metrics(minibatch=True, epoch=True)
        self.assertEqual(2, len(metrics))
        self.assertIsInstance(metrics, List)
        self.assertIsInstance(metrics[0], PluginMetric)
        self.assertIsInstance(metrics[1], PluginMetric)

class LossMetricTests(unittest.TestCase):
    def setUp(self):
        common_setups()

    def test_standalone_forgetting(self):
        uut = Loss()

        # Initial loss should be 0
        self.assertEqual(0.0, uut.result())

        loss = torch.as_tensor([0.0, 0.1, 0.2, 0.3, 0.4])  # Avg = 0.2
        uut.update(loss, 5)

        self.assertAlmostEqual(0.2, uut.result())

        loss = torch.as_tensor([0.1, 0.2, 0.3, 0.4, 0.5])  # Avg = 0.3
        uut.update(loss, 5)

        self.assertAlmostEqual(0.25, uut.result())

        # After-reset loss should be 0
        uut.reset()
        self.assertEqual(0.0, uut.result())

        # Check if handles 0 loss
        loss = torch.as_tensor([0.0, 0.0, 0.0])
        uut.update(loss, 3)

        self.assertEqual(0.0, uut.result())

        # Check handles losses with different reductions
        uut.reset()
        loss = torch.rand((5, 20))
        expected_mean = loss.mean().item()
        uut.update(loss, loss.shape[0])

        self.assertAlmostEqual(expected_mean, uut.result())

        # Check that the last call to result didn't change the value
        self.assertAlmostEqual(expected_mean, uut.result())

    def test_loss_helper(self):
        metrics = loss_metrics(minibatch=True, epoch=True)
        self.assertEqual(2, len(metrics))
        self.assertIsInstance(metrics, List)
        self.assertIsInstance(metrics[0], PluginMetric)
        self.assertIsInstance(metrics[1], PluginMetric)

class ConfusionMatrixMetricTests(unittest.TestCase):
    def setUp(self):
        common_setups()

    def test_standalone_cm_fixed_size(self):
        uut = ConfusionMatrix(num_classes=10)

        # Initial confusion matrix should be a 10x10 matrix filled with zeros
        self.assertTrue(torch.equal(torch.zeros((10, 10), dtype=torch.long),
                                    uut.result()))

        truth = torch.as_tensor([0, 5, 2, 1, 0])
        predicted = torch.as_tensor([2, 3, 2, 5, 0])
        uut.update(truth, predicted)

        expected = torch.zeros((10, 10), dtype=torch.long)
        expected[0][2] += 1
        expected[5][3] += 1
        expected[2][2] += 1
        expected[1][5] += 1
        expected[0][0] += 1

        self.assertTrue(torch.equal(expected, uut.result()))

        # After-reset matrix should be a 10x10 matrix filled with zeros
        uut.reset()
        self.assertTrue(torch.equal(torch.zeros((10, 10), dtype=torch.long),
                                    uut.result()))

        # Should throw exception when len(truth) != len(predicted)
        with self.assertRaises(ValueError):
            truth = torch.as_tensor([0, 0, 1, 0])
            predicted = torch.as_tensor([1, 1, 1])
            uut.update(truth, predicted)

        # Check that matrix didn't change after error
        self.assertTrue(torch.equal(torch.zeros((10, 10), dtype=torch.long),
                                    uut.result()))

        # Test logits / one-hot support
        uut.reset()

        # Test one-hot (truth)
        truth = torch.as_tensor([[0, 1], [1, 0], [1, 0], [1, 0]])
        predicted = torch.as_tensor([1, 1, 1, 1])
        uut.update(truth, predicted)

        expected = torch.zeros((10, 10), dtype=torch.long)
        expected[1][1] += 1
        expected[0][1] += 1
        expected[0][1] += 1
        expected[0][1] += 1

        self.assertTrue(torch.equal(expected, uut.result()))

        # Test logits (predictions)
        truth = torch.as_tensor([1, 1, 0, 0])
        predicted = torch.as_tensor(
            [[0.73, 0.1], [0.22, 0.33],
             [0.99, 0.01], [0.12, 0.11]])
        uut.update(truth, predicted)

        expected[1][0] += 1
        expected[1][1] += 1
        expected[0][0] += 1
        expected[0][0] += 1

        self.assertTrue(torch.equal(expected, uut.result()))

        # Test one-hot (truth) + logits (predictions)
        truth = torch.as_tensor([[1, 0], [1, 0], [0, 1], [0, 1]])
        predicted = torch.as_tensor(
            [[0.73, 0.1], [0.22, 0.33],
             [0.99, 0.01], [0.12, 0.11]])
        uut.update(truth, predicted)

        expected[0][0] += 1
        expected[0][1] += 1
        expected[1][0] += 1
        expected[1][0] += 1

        self.assertTrue(torch.equal(expected, uut.result()))

    def test_standalone_cm_error_handling(self):
        uut = ConfusionMatrix(num_classes=10)

        # Initial confusion matrix should be a 10x10 matrix filled with zeros
        self.assertTrue(torch.equal(torch.zeros((10, 10), dtype=torch.long),
                                    uut.result()))

        truth = torch.as_tensor([0, 5, 2, 1, 0])
        predicted = torch.as_tensor([2, 3, -1, 5, 0])

        with self.assertRaises(ValueError):
            uut.update(truth, predicted)

        truth = torch.as_tensor([0, 5, 2, 1, -5])
        predicted = torch.as_tensor([2, 3, 1, 5, 0])

        with self.assertRaises(ValueError):
            uut.update(truth, predicted)

class CpuUsageMetricTests(unittest.TestCase):
    def setUp(self):
        common_setups()

    def test_standalone_cpu_usage(self):
        uut = CPUUsage()

        # Assert result is 0 when created
        self.assertEqual(0.0, uut.result())

        # Base usage
        uut.update()
        for i in range(5):
            a = 0
            for j in range(1000000):
                a += 1
            uut.update()
            self.assertLessEqual(0.0, uut.result())
            time.sleep(0.3)

        # Assert reset actually resets
        uut.reset()
        self.assertEqual(0.0, uut.result())

        # Assert reset doesn't restart the tracking
        a = 0
        for j in range(1000000):
            a += 1
        time.sleep(0.3)
        uut.update()
        self.assertEqual(0.0, uut.result())

        # Assert result doesn't update the metric value
        a = 0
        for j in range(1000000):
            a += 1
        time.sleep(0.3)

        last_result = uut.result()
        self.assertLessEqual(0.0, last_result)

        a = 0
        for j in range(1000000):
            a += 1
        time.sleep(0.3)

        self.assertEqual(last_result, uut.result())

    def test_cpu_usage_helper(self):
        metrics = cpu_usage_metrics(epoch=True, experience=True)
        self.assertEqual(2, len(metrics))
        self.assertIsInstance(metrics, List)
        self.assertIsInstance(metrics[0], PluginMetric)
        self.assertIsInstance(metrics[1], PluginMetric)

class RamUsageMetricTests(unittest.TestCase):
    def setUp(self):
        common_setups()

    def test_standalone_cpu_usage(self):
        uut = MaxRAM()

        # Assert result is None when created
        self.assertEqual(0, uut.result())
        uut.start_thread()

        # Base usage
        last_result = uut.result()
        self.assertLessEqual(0.0, last_result)

        huge_blob = torch.rand(100000, dtype=torch.float64)

        self.assertLessEqual(last_result, uut.result())
        self.assertLessEqual(-1, huge_blob[0].item())  # Just to prevent removal

        # Assert reset actually resets
        uut.reset()
        self.assertEqual(0, uut.result())

class DiskUsageMetricTests(unittest.TestCase):
    def setUp(self):
        common_setups()

    def test_standalone_disk_usage(self):
        with tempfile.TemporaryDirectory() as tmp_dir:
            uut = DiskUsage(tmp_dir)

            # Assert result is None when created
            self.assertEqual(0, uut.result())

            # Base usage
            uut.update()
            self.assertLessEqual(0, uut.result())

            # Assert reset actually resets
            uut.reset()
            self.assertEqual(0, uut.result())

            uut.update()
            base_value = uut.result()
            with open(str(Path(tmp_dir) / 'blob_file.bin'), "wb") as f:
                f.write(os.urandom(512))

            # Shouldn't change between calls to update
            self.assertEqual(base_value, uut.result())

            uut.update()

            self.assertLessEqual(base_value, uut.result())

    def test_standalone_disk_usage_multi_dir(self):
        with tempfile.TemporaryDirectory() as tmp_dirA:
            with tempfile.TemporaryDirectory() as tmp_dirB:
                uut = DiskUsage([tmp_dirA, tmp_dirB])

                uut.update()
                base_value = uut.result()
                with open(str(Path(tmp_dirA) / 'blob_fileA.bin'), "wb") as f:
                    f.write(os.urandom(512))

                uut.update()
                self.assertLessEqual(base_value, uut.result())

                base_value = uut.result()

                with open(str(Path(tmp_dirB) / 'blob_fileB.bin'), "wb") as f:
                    f.write(os.urandom(512))

                uut.update()
                self.assertLessEqual(base_value, uut.result())

if __name__ == '__main__':
    unittest.main()
<commit_msg>Removed old tests<commit_after>""" Metrics Tests"""

import unittest

import torch

from avalanche.evaluation import PluginMetric
from avalanche.evaluation.metrics import Accuracy, Loss, ConfusionMatrix, \
    DiskUsage, MAC, CPUUsage, MaxGPU, MaxRAM, Mean, Sum, ElapsedTime, Forgetting
from avalanche.evaluation.metrics import accuracy_metrics, loss_metrics, StreamConfusionMatrix, \
    disk_usage_metrics, MAC_metrics, cpu_usage_metrics, gpu_usage_metrics, ram_usage_metrics, \
    timing_metrics

#################################
#################################
#### STANDALONE METRIC TEST #####
#################################
#################################

class GeneralMetricTests(unittest.TestCase):
    def setUp(self) -> None:
        self.batch_size = 3
        self.input_size = 10
        self.n_classes = 3
        self.out = torch.randn(self.batch_size, self.input_size)
        self.y = torch.randint(0, self.n_classes, (self.batch_size,))

    def test_accuracy(self):
        metric = Accuracy()
        self.assertEqual(metric.result(), 0)
        metric.update(self.out, self.y)
        self.assertLessEqual(metric.result(), 1)
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_loss(self):
        metric = Loss()
        self.assertEqual(metric.result(), 0)
        metric.update(torch.tensor(1.), self.batch_size)
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_cm(self):
        metric = ConfusionMatrix()
        cm = metric.result()
        self.assertTrue((cm == 0).all().item())
        metric.update(self.y, self.out)
        cm = metric.result()
        self.assertTrue((cm >= 0).all().item())
        metric.reset()
        cm = metric.result()
        self.assertTrue((cm == 0).all().item())

    def test_ram(self):
        metric = MaxRAM()
        self.assertEqual(metric.result(), 0)
        metric.start_thread()  # start thread
        self.assertGreaterEqual(metric.result(), 0)
        metric.stop_thread() # stop thread
        metric.reset()  # stop thread
        self.assertEqual(metric.result(), 0)

    def test_gpu(self):
        if torch.cuda.is_available():
            metric = MaxGPU(0)
            self.assertEqual(metric.result(), 0)
            metric.start_thread()  # start thread
            self.assertGreaterEqual(metric.result(), 0)
            metric.stop_thread() # stop thread
            metric.reset()  # stop thread
            self.assertEqual(metric.result(), 0)

    def test_cpu(self):
        metric = CPUUsage()
        self.assertEqual(metric.result(), 0)
        metric.update()
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_disk(self):
        metric = DiskUsage()
        self.assertEqual(metric.result(), 0)
        metric.update()
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_timing(self):
        metric = ElapsedTime()
        self.assertEqual(metric.result(), 0)
        metric.update()  # need two update calls
        self.assertEqual(metric.result(), 0)
        metric.update()
        self.assertGreaterEqual(metric.result(), 0)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_mac(self):
        model = torch.nn.Linear(self.input_size, 2)
        metric = MAC()
        self.assertEqual(metric.result(), 0)
        metric.update(model, self.out)
        self.assertGreaterEqual(metric.result(), 0)

    def test_mean(self):
        metric = Mean()
        self.assertEqual(metric.result(), 0)
        metric.update(0.1, 1)
        self.assertEqual(metric.result(), 0.1)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_sum(self):
        metric = Sum()
        self.assertEqual(metric.result(), 0)
        metric.update(5)
        self.assertEqual(metric.result(), 5)
        metric.reset()
        self.assertEqual(metric.result(), 0)

    def test_forgetting(self):
        metric = Forgetting()
        f = metric.result()
        self.assertEqual(f, {})
        f = metric.result(k=0)
        self.assertIsNone(f)
        metric.update(0, 1, initial=True)
        f = metric.result(k=0)
        self.assertIsNone(f)
        metric.update(0, 0.4)
        f = metric.result(k=0)
        self.assertEqual(f, 0.6)
        metric.reset()
        self.assertEqual(metric.result(), {})

#################################
#################################
###### PLUGIN METRIC TEST #######
#################################
#################################

if __name__ == '__main__':
    unittest.main()
<|endoftext|>