<commit_before>import torch
from .constraints import positive
from torch import nn as nn
from torch.nn import functional as F
from torch.nn.modules.utils import _pair
# from .module import Module
from torch.nn import Parameter

class Offset(nn.Module):
    def __init__(self, offset=1):
        super().__init__()
        self.offset = offset

    def forward(self, x):
        return x + self.offset

class Elu1(nn.Module):

    def forward(self, x):
        return F.elu(x) + 1.

class Conv2dPad(nn.Conv2d):
    """
    Padded Conv2d layer. Pads with reflect by default.
    """

    def __init__(self, in_channels, out_channels, kernel_size, pad=None, mode='reflect', **kwargs):
        assert 'padding' not in kwargs, 'You need to use `pad` not `padding`'
        self.padding = _pair(0)
        super().__init__(in_channels, out_channels, kernel_size, padding=0, **kwargs)
        self.mode = mode

        if isinstance(pad, tuple) or pad is None:
            self.pad = pad
        else:
            self.pad = 4 * (pad,)

    def _pad(self, input):
        if self.pad is not None and self.pad != 4 * (0,):
            input = F.pad(input, mode=self.mode, pad=self.pad)
        return input

    def forward(self, input):
        input = self._pad(input)
        return F.conv2d(input, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)

class BiasBatchNorm2d(nn.BatchNorm2d):
    def __init__(self, features, **kwargs):
        super().__init__(features, affine=False, **kwargs)
        bias = nn.Parameter(torch.FloatTensor(1, features, 1, 1))
        self.register_parameter('bias', bias)

    def forward(self, input):
        self._check_input_dim(input)
        return F.batch_norm(input, self.running_mean, self.running_var, weight=None, bias=None,
                            training=self.training, momentum=self.momentum, eps=self.eps)
        N, _, w, h = input.size()
        return input + self.bias.repeat(N, 1, w, h)

class SpatialXFeatureLinear(nn.Module):
    """
    Factorized fully connected layer. Weights are a sum of outer products between a spatial filter and a feature vector.
    """

    def __init__(self, in_shape, outdims, bias=True, normalize=True, positive=False, spatial=None):
        super().__init__()
        self.in_shape = in_shape
        c, w, h = self.in_shape
        self.outdims = outdims
        self.normalize = normalize
        self.positive = positive

        self.spatial = Parameter(torch.Tensor(self.outdims, 1, w, h)) if spatial is None else spatial
        self.features = Parameter(torch.Tensor(self.outdims, c, 1, 1))

        if bias:
            bias = Parameter(torch.Tensor(self.outdims))
            self.register_parameter('bias', bias)
        else:
            self.register_parameter('bias', None)
        self.initialize()

    @property
    def normalized_spatial(self):
        if self.positive:
            positive(self.spatial)
        if self.normalize:
            weight = self.spatial / (self.spatial.pow(2).sum(2).sum(3).sqrt().expand_as(self.spatial) + 1e-6)
        else:
            weight = self.spatial
        return weight

    @property
    def weight(self):
        c, w, h = self.in_shape
        n = self.outdims
        weight = self.spatial.expand(n, c, w, h) * self.features.expand(n, c, w, h)
        weight = weight.view(self.outdims, -1)
        return weight

    def initialize(self, init_noise=1e-3):
        self.spatial.data.normal_(0, init_noise)
        self.features.data.normal_(0, init_noise)
        if self.bias is not None:
            self.bias.data.fill_(0)

    def forward(self, x):
        N = x.size(0)
        y = x.view(N, -1) @ self.weight.t()
        if self.bias is not None:
            y = y + self.bias.expand_as(y)
        return y

    def __repr__(self):
        return ('spatial positive ' if self.positive else '') + \
               ('normalized ' if self.normalize else '') + \
               self.__class__.__name__ + \
               ' (' + '{} x {} x {}'.format(*self.in_shape) + ' -> ' + str(
            self.outdims) + ')'

class WidthXHeightXFeatureLinear(nn.Module):
    """
    Factorized fully connected layer. Weights are a sum of outer products between three vectors over width,
    height and spatial.
    """

    def __init__(self, in_shape, outdims, components=1, bias=True, normalize=True, positive=False, width=None,
                 height=None, eps=1e-6):
        super().__init__()
        self.in_shape = in_shape
        self.eps = eps
        c, w, h = self.in_shape
        self.outdims = outdims
        self.normalize = normalize
        self.positive = positive
        self.components = components

        self.width = Parameter(torch.Tensor(self.outdims, 1, w, 1, components)) if width is None else width
        self.height = Parameter(torch.Tensor(self.outdims, 1, 1, h, components)) if height is None else height
        self.features = Parameter(torch.Tensor(self.outdims, c, 1, 1))
        assert self.width.size(4) == self.height.size(4), 'The number of components in width and height do not agree'
        self.components = self.width.size(4)
        if bias:
            bias = Parameter(torch.Tensor(self.outdims))
            self.register_parameter('bias', bias)
        else:
            self.register_parameter('bias', None)
        self.initialize()

    def initialize(self, init_noise=1e-3):
        self.width.data.normal_(0, init_noise)
        self.height.data.normal_(0, init_noise)
        self.features.data.normal_(0, init_noise)
        if self.bias is not None:
            self.bias.data.fill_(0)

    @property
    def normalized_width(self):
        if self.positive:
            positive(self.width)
        if self.normalize:
            return self.width / (self.width.pow(2).sum(2).sqrt().expand_as(self.width) + self.eps)
        else:
            return self.width

    @property
    def normalized_height(self):
        c, w, h = self.in_shape
        if self.positive:
            positive(self.height)
        if self.normalize:
            return self.height / (self.height.pow(2).sum(3).sqrt().expand_as(self.height) + self.eps)
        else:
            return self.height

    @property
    def spatial(self):
        c, w, h = self.in_shape
        n, comp = self.outdims, self.components
        weight = self.normalized_width.expand(n, 1, w, h, comp) \
                 * self.normalized_height.expand(n, 1, w, h, comp)
        weight = weight.sum(4).view(n, 1, w, h)
        return weight

    @property
    def weight(self):
        c, w, h = self.in_shape
        n, comp = self.outdims, self.components
        weight = self.spatial.expand(n, c, w, h) * self.features.expand(n, c, w, h)
        weight = weight.view(self.outdims, -1)
        return weight

    @property
    def basis(self):
        c, w, h = self.in_shape
        return self.weight.view(-1, c, w, h).data.cpu().numpy()

    def forward(self, x):

        N = x.size(0)
        y = x.view(N, -1) @ self.weight.t()
        if self.bias is not None:
            y = y + self.bias.expand_as(y)
        return y

    def __repr__(self):
        return ('spatial positive ' if self.positive else '') + \
               ('normalized ' if self.normalize else '') + \
               self.__class__.__name__ + \
               ' (' + '{} x {} x {}'.format(*self.in_shape) + ' -> ' + str(self.outdims) + ') spatial rank {}'.format(
            self.components)

class BiasBatchNorm2d(nn.BatchNorm2d):
    def __init__(self, features, **kwargs):
        kwargs['affine'] = False
        super().__init__(features, **kwargs)
        self.bias = nn.Parameter(torch.FloatTensor(1, features, 1, 1))
        self.initialize()

    def initialize(self, init_noise=None):
        self.bias.data.fill_(0)

    def forward(self, input):
        output = super().forward(input)
        return output + self.bias.expand_as(output)
<commit_msg>Remove duplicate BiasBatchNorm2d<commit_after>import torch
from .constraints import positive
from torch import nn as nn
from torch.nn import functional as F
from torch.nn.modules.utils import _pair
# from .module import Module
from torch.nn import Parameter

class Offset(nn.Module):
    def __init__(self, offset=1):
        super().__init__()
        self.offset = offset

    def forward(self, x):
        return x + self.offset

class Elu1(nn.Module):

    def forward(self, x):
        return F.elu(x) + 1.

class Conv2dPad(nn.Conv2d):
    """
    Padded Conv2d layer. Pads with reflect by default.
    """

    def __init__(self, in_channels, out_channels, kernel_size, pad=None, mode='reflect', **kwargs):
        assert 'padding' not in kwargs, 'You need to use `pad` not `padding`'
        self.padding = _pair(0)
        super().__init__(in_channels, out_channels, kernel_size, padding=0, **kwargs)
        self.mode = mode

        if isinstance(pad, tuple) or pad is None:
            self.pad = pad
        else:
            self.pad = 4 * (pad,)

    def _pad(self, input):
        if self.pad is not None and self.pad != 4 * (0,):
            input = F.pad(input, mode=self.mode, pad=self.pad)
        return input

    def forward(self, input):
        input = self._pad(input)
        return F.conv2d(input, self.weight, self.bias, self.stride, self.padding, self.dilation, self.groups)

class SpatialXFeatureLinear(nn.Module):
    """
    Factorized fully connected layer. Weights are a sum of outer products between a spatial filter and a feature vector.
    """

    def __init__(self, in_shape, outdims, bias=True, normalize=True, positive=False, spatial=None):
        super().__init__()
        self.in_shape = in_shape
        c, w, h = self.in_shape
        self.outdims = outdims
        self.normalize = normalize
        self.positive = positive

        self.spatial = Parameter(torch.Tensor(self.outdims, 1, w, h)) if spatial is None else spatial
        self.features = Parameter(torch.Tensor(self.outdims, c, 1, 1))

        if bias:
            bias = Parameter(torch.Tensor(self.outdims))
            self.register_parameter('bias', bias)
        else:
            self.register_parameter('bias', None)
        self.initialize()

    @property
    def normalized_spatial(self):
        if self.positive:
            positive(self.spatial)
        if self.normalize:
            weight = self.spatial / (self.spatial.pow(2).sum(2).sum(3).sqrt().expand_as(self.spatial) + 1e-6)
        else:
            weight = self.spatial
        return weight

    @property
    def weight(self):
        c, w, h = self.in_shape
        n = self.outdims
        weight = self.spatial.expand(n, c, w, h) * self.features.expand(n, c, w, h)
        weight = weight.view(self.outdims, -1)
        return weight

    def initialize(self, init_noise=1e-3):
        self.spatial.data.normal_(0, init_noise)
        self.features.data.normal_(0, init_noise)
        if self.bias is not None:
            self.bias.data.fill_(0)

    def forward(self, x):
        N = x.size(0)
        y = x.view(N, -1) @ self.weight.t()
        if self.bias is not None:
            y = y + self.bias.expand_as(y)
        return y

    def __repr__(self):
        return ('spatial positive ' if self.positive else '') + \
               ('normalized ' if self.normalize else '') + \
               self.__class__.__name__ + \
               ' (' + '{} x {} x {}'.format(*self.in_shape) + ' -> ' + str(
            self.outdims) + ')'

class WidthXHeightXFeatureLinear(nn.Module):
    """
    Factorized fully connected layer. Weights are a sum of outer products between three vectors over width,
    height and spatial.
    """

    def __init__(self, in_shape, outdims, components=1, bias=True, normalize=True, positive=False, width=None,
                 height=None, eps=1e-6):
        super().__init__()
        self.in_shape = in_shape
        self.eps = eps
        c, w, h = self.in_shape
        self.outdims = outdims
        self.normalize = normalize
        self.positive = positive
        self.components = components

        self.width = Parameter(torch.Tensor(self.outdims, 1, w, 1, components)) if width is None else width
        self.height = Parameter(torch.Tensor(self.outdims, 1, 1, h, components)) if height is None else height
        self.features = Parameter(torch.Tensor(self.outdims, c, 1, 1))
        assert self.width.size(4) == self.height.size(4), 'The number of components in width and height do not agree'
        self.components = self.width.size(4)
        if bias:
            bias = Parameter(torch.Tensor(self.outdims))
            self.register_parameter('bias', bias)
        else:
            self.register_parameter('bias', None)
        self.initialize()

    def initialize(self, init_noise=1e-3):
        self.width.data.normal_(0, init_noise)
        self.height.data.normal_(0, init_noise)
        self.features.data.normal_(0, init_noise)
        if self.bias is not None:
            self.bias.data.fill_(0)

    @property
    def normalized_width(self):
        if self.positive:
            positive(self.width)
        if self.normalize:
            return self.width / (self.width.pow(2).sum(2).sqrt().expand_as(self.width) + self.eps)
        else:
            return self.width

    @property
    def normalized_height(self):
        c, w, h = self.in_shape
        if self.positive:
            positive(self.height)
        if self.normalize:
            return self.height / (self.height.pow(2).sum(3).sqrt().expand_as(self.height) + self.eps)
        else:
            return self.height

    @property
    def spatial(self):
        c, w, h = self.in_shape
        n, comp = self.outdims, self.components
        weight = self.normalized_width.expand(n, 1, w, h, comp) \
                 * self.normalized_height.expand(n, 1, w, h, comp)
        weight = weight.sum(4).view(n, 1, w, h)
        return weight

    @property
    def weight(self):
        c, w, h = self.in_shape
        n, comp = self.outdims, self.components
        weight = self.spatial.expand(n, c, w, h) * self.features.expand(n, c, w, h)
        weight = weight.view(self.outdims, -1)
        return weight

    @property
    def basis(self):
        c, w, h = self.in_shape
        return self.weight.view(-1, c, w, h).data.cpu().numpy()

    def forward(self, x):

        N = x.size(0)
        y = x.view(N, -1) @ self.weight.t()
        if self.bias is not None:
            y = y + self.bias.expand_as(y)
        return y

    def __repr__(self):
        return ('spatial positive ' if self.positive else '') + \
               ('normalized ' if self.normalize else '') + \
               self.__class__.__name__ + \
               ' (' + '{} x {} x {}'.format(*self.in_shape) + ' -> ' + str(self.outdims) + ') spatial rank {}'.format(
            self.components)

class BiasBatchNorm2d(nn.BatchNorm2d):
    def __init__(self, features, **kwargs):
        kwargs['affine'] = False
        super().__init__(features, **kwargs)
        self.bias = nn.Parameter(torch.FloatTensor(1, features, 1, 1))
        self.initialize()

    def initialize(self, init_noise=None):
        self.bias.data.fill_(0)

    def forward(self, input):
        output = super().forward(input)
        return output + self.bias.expand_as(output)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

from __future__ import print_function

from .._protos.public.modeldb.versioning import VersioningService_pb2 as _VersioningService

class Diff(object):
    def __init__(self, diffs):
        # TODO: order so that additions come after removal?
        self._diffs = diffs

    def __repr__(self):
        return '\n'.join(
            '\n'.join((
                "{} {}".format(
                    _VersioningService.DiffStatusEnum.DiffStatus.Name(diff.status),
                    '/'.join(diff.location),
                ),
                getattr(diff, diff.WhichOneof('content')).__repr__(),
            ))
            for diff
            in self._diffs
        )
<commit_msg>Update DiffStatusEnum location (#431)<commit_after># -*- coding: utf-8 -*-

from __future__ import print_function

from .._protos.public.modeldb.versioning import Enums_pb2 as _VersioningEnumsProtos

class Diff(object):
    def __init__(self, diffs):
        # TODO: order so that additions come after removal?
        self._diffs = diffs

    def __repr__(self):
        return '\n'.join(
            '\n'.join((
                "{} {}".format(
                    _VersioningEnumsProtos.DiffStatusEnum.DiffStatus.Name(diff.status),
                    '/'.join(diff.location),
                ),
                getattr(diff, diff.WhichOneof('content')).__repr__(),
            ))
            for diff
            in self._diffs
        )
<|endoftext|>


<commit_before>'''
A scenario discovery oriented implementation of CART. It essentially is a
wrapper around scikit-learn's version of CART.

'''
from __future__ import (absolute_import, print_function, division,
                        unicode_literals)

import io
import math

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import numpy as np
import pandas as pd
import seaborn as sns
import six
from sklearn import tree
from sklearn.externals.six import StringIO

from . import scenario_discovery_util as sdutil

# Created on May 22, 2015
#
# .. codeauthor:: jhkwakkel <j.h.kwakkel (at) tudelft (dot) nl>

__all__ = ['setup_cart',
           'CART']

def setup_cart(results, classify, incl_unc=[], mass_min=0.05):
    """helper function for performing cart in combination with data
    generated by the workbench.

    Parameters
    ----------
    results : tuple of structured array and dict with numpy arrays
              the return from :meth:`perform_experiments`.
    classify : string, function or callable
               either a string denoting the outcome of interest to
               use or a function.
    incl_unc : list of strings
    mass_min : float

    Raises
    ------
    TypeError
        if classify is not a string or a callable.

    """
    x, y, mode = sdutil._setup(results, classify, incl_unc)

    return CART(x, y, mass_min, mode=mode)

class CART(sdutil.OutputFormatterMixin):
    '''CART algorithm

    can be used in a manner similar to PRIM. It provides access
    to the underlying tree, but it can also show the boxes described by the
    tree in a table or graph form similar to prim.

    Parameters
    ----------
    x : recarray
    y : ndarray
    mass_min : float, optional
               a value between 0 and 1 indicating the minimum fraction
               of data points in a terminal leaf. Defaults to 0.05,
               identical to prim.
    mode : {BINARY, CLASSIFICATION, REGRESSION}
           indicates the mode in which CART is used. Binary indicates
           binary classification, classification is multiclass, and regression
           is regression.

    Attributes
    ----------
    boxes : list
            list of recarray box lims
    stats : list
            list of dicts with stats

    Notes
    -----
    This class is a wrapper around scikit-learn's CART algorithm. It provides
    an interface to CART that is more oriented towards scenario discovery, and
    shared some methods with PRIM

    See also
    --------
    :mod:`prim`

    '''

    sep = '?!?'

    def __init__(self, x, y, mass_min=0.05, mode=sdutil.BINARY):
        ''' init

        '''
        x = x.drop(["scenario"], axis=1)

        self.x = x
        self.y = y
        self.mass_min = mass_min
        self.mode = mode

        # we need to transform the structured array to a ndarray
        # we use dummy variables for each category in case of categorical
        # variables. Integers are treated as floats
        dummies = pd.get_dummies(self.x)

        self.feature_names = dummies.columns.values.tolist()
        self._x = dummies.values
        self._boxes = None
        self._stats = None

    @property
    def boxes(self):
        assert self.clf

        if self._boxes:
            return self._boxes

        # based on
        # http://stackoverflow.com/questions/20224526/how-to-extract-the-
        # decision-rules-from-scikit-learn-decision-tree
        assert self.clf

        left = self.clf.tree_.children_left
        right = self.clf.tree_.children_right
        threshold = self.clf.tree_.threshold
        features = [self.feature_names[i] for i in self.clf.tree_.feature]

        # get ids of leaf nodes
        leafs = np.argwhere(left == -1)[:, 0]

        def recurse(left, right, child, lineage=None):
            if lineage is None:
                # lineage = [self.clf.tree_.value[child]]
                lineage = []

            if child in left:
                parent = np.where(left == child)[0].item()
                split = 'l'
            else:
                parent = np.where(right == child)[0].item()
                split = 'r'

            lineage.append((parent, split, threshold[parent],
                            features[parent]))

            if parent == 0:
                lineage.reverse()
                return lineage
            else:
                return recurse(left, right, parent, lineage)

        box_init = sdutil._make_box(self.x)
        boxes = []
        for leaf in leafs:
            branch = recurse(left, right, leaf)
            box = box_init.copy()
            for node in branch:
                direction = node[1]
                value = node[2]
                unc = node[3]

                if direction == 'l':
                    try:
                        box.loc[1, unc] = value
                    except ValueError:
                        unc, cat = unc.split(self.sep)
                        cats = list(box[unc][0])
                        cats = [str(cat) for cat in cats]
                        cats.pop(cats.index(str(cat)))
                        box.loc[:, unc] = set(cats)
                else:
                    try:
                        if box[unc].dtype == np.int32:
                            value = math.ceil(value)
                        box.loc[0, unc] = value
                    except (ValueError, KeyError):
                        # we are in the right hand branch, so
                        # the category is included
                        pass

            boxes.append(box)
        self._boxes = boxes
        return self._boxes

    @property
    def stats(self):
        if self._stats:
            return self._stats

        boxes = self.boxes

        box_init = sdutil._make_box(self.x)

        self._stats = []
        for box in boxes:
            boxstats = self._boxstat_methods[self.mode](self, box, box_init)
            self._stats.append(boxstats)
        return self._stats

    def _binary_stats(self, box, box_init):
        logical = sdutil._in_box(self.x, box)

        y_in_box = self.y[logical]
        box_coi = np.sum(y_in_box)

        boxstats = {'coverage': box_coi/np.sum(self.y),
                    'density': box_coi/y_in_box.shape[0],
                    'res dim': sdutil._determine_nr_restricted_dims(box,
                                                                    box_init),
                    'mass': y_in_box.shape[0]/self.y.shape[0]}
        return boxstats

    def _regression_stats(self, box, box_init):
        logical = sdutil._in_box(self.x, box)

        y_in_box = self.y[logical]

        boxstats = {'mean': np.mean(y_in_box),
                    'mass': y_in_box.shape[0]/self.y.shape[0],
                    'res dim': sdutil._determine_nr_restricted_dims(box,
                                                                    box_init)}
        return boxstats

    def _classification_stats(self, box, box_init):
        logical = sdutil._in_box(self.x, box)

        y_in_box = self.y[logical]
        classes = set(self.y)
        classes = list(classes)
        classes.sort()

        counts = [y_in_box[y_in_box == ci].shape[0] for ci in classes]

        total_gini = 0
        for count in counts:
            total_gini += (count/y_in_box.shape[0])**2
        gini = 1 - total_gini

        boxstats = {'gini': gini,
                    'mass': y_in_box.shape[0]/self.y.shape[0],
                    'box_composition': counts,
                    'res dim': sdutil._determine_nr_restricted_dims(box,
                                                                    box_init)}

        return boxstats

    _boxstat_methods = {sdutil.BINARY: _binary_stats,
                        sdutil.REGRESSION: _regression_stats,
                        sdutil.CLASSIFICATION: _classification_stats}

    def build_tree(self):
        '''train CART on the data'''
        min_samples = int(self.mass_min*self.x.shape[0])

        if self.mode == sdutil.REGRESSION:
            self.clf = tree.DecisionTreeRegressor(min_samples_leaf=min_samples)
        else:
            self.clf = tree.DecisionTreeClassifier(
                min_samples_leaf=min_samples)
        self.clf.fit(self._x, self.y)

    def show_tree(self, mplfig=True):
        '''return a png of the tree

        Parameters
        ----------
        mplfig : bool, optional
                 if true (default) returns a matplotlib figure with the tree,
                 otherwise, it returns a png as bytes

        '''
        assert self.clf
        try:
            import pydotplus as pydot
        except ImportError:
            import pydot  # dirty hack for read the docs

        dot_data = StringIO()
        tree.export_graphviz(self.clf, out_file=dot_data,
                             feature_names=self.feature_names)
        dot_data = dot_data.getvalue()  # .encode('ascii') # @UndefinedVariable
        graph = pydot.graph_from_dot_data(dot_data)
        img = graph.create_png()

        if mplfig:
            fig, ax = plt.subplots()
            ax.imshow(mpimg.imread(io.BytesIO(img)))
            ax.axis('off')
            return fig

        return img

# if __name__ == '__main__':
#     from test import test_utilities
#     import matplotlib.pyplot as plt
#
#     ema_logging.log_to_stderr(ema_logging.INFO)
#
#     def scarcity_classify(outcomes):
#         outcome = outcomes['relative market price']
#         change = np.abs(outcome[:, 1::]-outcome[:, 0:-1])
#
#         neg_change = np.min(change, axis=1)
#         pos_change = np.max(change, axis=1)
#
#         logical = (neg_change > -0.6) & (pos_change > 0.6)
#
#         classes = np.zeros(outcome.shape[0])
#         classes[logical] = 1
#
#         return classes
#
#     results = test_utilities.load_scarcity_data()
#
#     cart = setup_cart(results, scarcity_classify)
#     cart.build_tree()
#
#     print(cart.boxes_to_dataframe())
#     print(cart.stats_to_dataframe())
#     cart.display_boxes(together=True)
#
#     img = cart.show_tree()
#
#     import matplotlib.pyplot as plt
#     import matplotlib.image as mpimg
#
#     # treat the dot output string as an image file
#     sio = StringIO()
#     sio.write(img)
#     sio.seek(0)
#     img = mpimg.imread(sio)
#
#     # plot the image
#     imgplot = plt.imshow(img, aspect='equal')
#
#     plt.show()
<commit_msg>Revert "Revert "separator for dummies and updated docs""<commit_after>'''
A scenario discovery oriented implementation of CART. It essentially is a
wrapper around scikit-learn's version of CART.

'''
from __future__ import (absolute_import, print_function, division,
                        unicode_literals)

import io
import math

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

import numpy as np
import pandas as pd

import six
from sklearn import tree
from sklearn.externals.six import StringIO

from . import scenario_discovery_util as sdutil

# Created on May 22, 2015
#
# .. codeauthor:: jhkwakkel <j.h.kwakkel (at) tudelft (dot) nl>

__all__ = ['setup_cart',
           'CART']

def setup_cart(results, classify, incl_unc=[], mass_min=0.05):
    """helper function for performing cart in combination with data
    generated by the workbench.

    Parameters
    ----------
    results : tuple of DataFrame and dict with numpy arrays
              the return from :meth:`perform_experiments`.
    classify : string, function or callable
               either a string denoting the outcome of interest to
               use or a function.
    incl_unc : list of strings
    mass_min : float

    Raises
    ------
    TypeError
        if classify is not a string or a callable.

    """
    x, y, mode = sdutil._setup(results, classify, incl_unc)

    return CART(x, y, mass_min, mode=mode)

class CART(sdutil.OutputFormatterMixin):
    '''CART algorithm

    can be used in a manner similar to PRIM. It provides access
    to the underlying tree, but it can also show the boxes described by the
    tree in a table or graph form similar to prim.

    Parameters
    ----------
    x : DataFrame
    y : 1D ndarray
    mass_min : float, optional
               a value between 0 and 1 indicating the minimum fraction
               of data points in a terminal leaf. Defaults to 0.05,
               identical to prim.
    mode : {BINARY, CLASSIFICATION, REGRESSION}
           indicates the mode in which CART is used. Binary indicates
           binary classification, classification is multiclass, and regression
           is regression.

    Attributes
    ----------
    boxes : list
            list of DataFrame box lims
    stats : list
            list of dicts with stats

    Notes
    -----
    This class is a wrapper around scikit-learn's CART algorithm. It provides
    an interface to CART that is more oriented towards scenario discovery, and
    shared some methods with PRIM

    See also
    --------
    :mod:`prim`

    '''

    sep = '!?!'

    def __init__(self, x, y, mass_min=0.05, mode=sdutil.BINARY):
        ''' init

        '''
        x = x.drop(["scenario"], axis=1)

        self.x = x
        self.y = y
        self.mass_min = mass_min
        self.mode = mode

        # we need to transform the structured array to a ndarray
        # we use dummy variables for each category in case of categorical
        # variables. Integers are treated as floats
        dummies = pd.get_dummies(self.x, prefix_sep=self.sep)

        self.feature_names = dummies.columns.values.tolist()
        self._x = dummies.values
        self._boxes = None
        self._stats = None

    @property
    def boxes(self):
        assert self.clf

        if self._boxes:
            return self._boxes

        # based on
        # http://stackoverflow.com/questions/20224526/how-to-extract-the-
        # decision-rules-from-scikit-learn-decision-tree
        assert self.clf

        left = self.clf.tree_.children_left
        right = self.clf.tree_.children_right
        threshold = self.clf.tree_.threshold
        features = [self.feature_names[i] for i in self.clf.tree_.feature]

        # get ids of leaf nodes
        leafs = np.argwhere(left == -1)[:, 0]

        def recurse(left, right, child, lineage=None):
            if lineage is None:
                # lineage = [self.clf.tree_.value[child]]
                lineage = []

            if child in left:
                parent = np.where(left == child)[0].item()
                split = 'l'
            else:
                parent = np.where(right == child)[0].item()
                split = 'r'

            lineage.append((parent, split, threshold[parent],
                            features[parent]))

            if parent == 0:
                lineage.reverse()
                return lineage
            else:
                return recurse(left, right, parent, lineage)

        box_init = sdutil._make_box(self.x)
        boxes = []
        for leaf in leafs:
            branch = recurse(left, right, leaf)
            box = box_init.copy()
            for node in branch:
                direction = node[1]
                value = node[2]
                unc = node[3]

                if direction == 'l':
                    try:
                        box.loc[1, unc] = value
                    except ValueError:
                        unc, cat = unc.split(self.sep)
                        cats = list(box[unc][0])
                        cats = [str(cat) for cat in cats]
                        cats.pop(cats.index(str(cat)))
                        box.loc[:, unc] = set(cats)
                else:
                    try:
                        if box[unc].dtype == np.int32:
                            value = math.ceil(value)
                        box.loc[0, unc] = value
                    except (ValueError, KeyError):
                        # we are in the right hand branch, so
                        # the category is included
                        pass

            boxes.append(box)
        self._boxes = boxes
        return self._boxes

    @property
    def stats(self):
        if self._stats:
            return self._stats

        boxes = self.boxes

        box_init = sdutil._make_box(self.x)

        self._stats = []
        for box in boxes:
            boxstats = self._boxstat_methods[self.mode](self, box, box_init)
            self._stats.append(boxstats)
        return self._stats

    def _binary_stats(self, box, box_init):
        logical = sdutil._in_box(self.x, box)

        y_in_box = self.y[logical]
        box_coi = np.sum(y_in_box)

        boxstats = {'coverage': box_coi/np.sum(self.y),
                    'density': box_coi/y_in_box.shape[0],
                    'res dim': sdutil._determine_nr_restricted_dims(box,
                                                                    box_init),
                    'mass': y_in_box.shape[0]/self.y.shape[0]}
        return boxstats

    def _regression_stats(self, box, box_init):
        logical = sdutil._in_box(self.x, box)

        y_in_box = self.y[logical]

        boxstats = {'mean': np.mean(y_in_box),
                    'mass': y_in_box.shape[0]/self.y.shape[0],
                    'res dim': sdutil._determine_nr_restricted_dims(box,
                                                                    box_init)}
        return boxstats

    def _classification_stats(self, box, box_init):
        logical = sdutil._in_box(self.x, box)

        y_in_box = self.y[logical]
        classes = set(self.y)
        classes = list(classes)
        classes.sort()

        counts = [y_in_box[y_in_box == ci].shape[0] for ci in classes]

        total_gini = 0
        for count in counts:
            total_gini += (count/y_in_box.shape[0])**2
        gini = 1 - total_gini

        boxstats = {'gini': gini,
                    'mass': y_in_box.shape[0]/self.y.shape[0],
                    'box_composition': counts,
                    'res dim': sdutil._determine_nr_restricted_dims(box,
                                                                    box_init)}

        return boxstats

    _boxstat_methods = {sdutil.BINARY: _binary_stats,
                        sdutil.REGRESSION: _regression_stats,
                        sdutil.CLASSIFICATION: _classification_stats}

    def build_tree(self):
        '''train CART on the data'''
        min_samples = int(self.mass_min*self.x.shape[0])

        if self.mode == sdutil.REGRESSION:
            self.clf = tree.DecisionTreeRegressor(min_samples_leaf=min_samples)
        else:
            self.clf = tree.DecisionTreeClassifier(
                min_samples_leaf=min_samples)
        self.clf.fit(self._x, self.y)

    def show_tree(self, mplfig=True):
        '''return a png of the tree

        Parameters
        ----------
        mplfig : bool, optional
                 if true (default) returns a matplotlib figure with the tree,
                 otherwise, it returns a png as bytes

        '''
        assert self.clf
        try:
            import pydotplus as pydot
        except ImportError:
            import pydot  # dirty hack for read the docs

        dot_data = StringIO()
        tree.export_graphviz(self.clf, out_file=dot_data,
                             feature_names=self.feature_names)
        dot_data = dot_data.getvalue()  # .encode('ascii') # @UndefinedVariable
        graph = pydot.graph_from_dot_data(dot_data)
        img = graph.create_png()

        if mplfig:
            fig, ax = plt.subplots()
            ax.imshow(mpimg.imread(io.BytesIO(img)))
            ax.axis('off')
            return fig

        return img

# if __name__ == '__main__':
#     from test import test_utilities
#     import matplotlib.pyplot as plt
#
#     ema_logging.log_to_stderr(ema_logging.INFO)
#
#     def scarcity_classify(outcomes):
#         outcome = outcomes['relative market price']
#         change = np.abs(outcome[:, 1::]-outcome[:, 0:-1])
#
#         neg_change = np.min(change, axis=1)
#         pos_change = np.max(change, axis=1)
#
#         logical = (neg_change > -0.6) & (pos_change > 0.6)
#
#         classes = np.zeros(outcome.shape[0])
#         classes[logical] = 1
#
#         return classes
#
#     results = test_utilities.load_scarcity_data()
#
#     cart = setup_cart(results, scarcity_classify)
#     cart.build_tree()
#
#     print(cart.boxes_to_dataframe())
#     print(cart.stats_to_dataframe())
#     cart.display_boxes(together=True)
#
#     img = cart.show_tree()
#
#     import matplotlib.pyplot as plt
#     import matplotlib.image as mpimg
#
#     # treat the dot output string as an image file
#     sio = StringIO()
#     sio.write(img)
#     sio.seek(0)
#     img = mpimg.imread(sio)
#
#     # plot the image
#     imgplot = plt.imshow(img, aspect='equal')
#
#     plt.show()
<|endoftext|>


<commit_before>"""Leetcode 101. Symmetric Tree
Easy

URL: https://leetcode.com/problems/symmetric-tree/

Given a binary tree, check whether it is a mirror of itself
(ie, symmetric around its center).

For example, this binary tree [1,2,2,3,4,4,3] is symmetric:
    1
   / \
  2   2
 / \ / \
3  4 4  3

But the following [1,2,2,null,3,null,3] is not:
    1
   / \
  2   2
   \   \
   3    3

Note:
Bonus points if you could solve it both recursively and iteratively.
"""

# Definition for a binary tree node.
class TreeNode(object):
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

class SolutionRecur(object):
    def isMirror(self, left, right):
        # Check left & right nodes's symmetry.
        if not left and not right:
            return True

        if not left or not right:
            return False

        # Check left & right nodes's values.
        if left.val != right.val:
            return False

        # Check outside & inside pairs.
        out_pair = self.isMirror(left.left, right.right)
        in_pair = self.isMirror(left.right, right.left)
        return out_pair and in_pair

    def isSymmetric(self, root):
        """
        :type root: TreeNode
        :rtype: bool

        Time complexity: O(n).
        Space complexity: O(logn) for balanced tree; O(n) for single sided tree.
        """
        if not root:
            return True

        return self.isMirror(root.left, root.right)

class SolutionIter(object):
    def isSymmetric(self, root):
        """
        :type root: TreeNode
        :rtype: bool

        Time complexity: O(n).
        Space complexity: O(logn) for balanced tree; O(n) for single sided tree.
        """
        if not root:
            return True

        queue = [(root.left, root.right)]

        while queue:
            left, right = queue.pop()

            # Check left & right nodes's symmetry.
            if not left and not right:
                continue

            if not left or not right:
                return False

            # Check left & right nodes's values.
            if left.val != right.val:
                return False

            # Check outside & inside pairs.
            queue.insert(0, (left.left, right.right))
            queue.insert(0, (left.right, right.left))

        return True

def main():
    # Output: True.
    #     1
    #    / \
    #   2   2
    #  / \ / \
    # 3  4 4  3
    root = TreeNode(1)
    root.left = TreeNode(2)
    root.right = TreeNode(2)
    root.left.left = TreeNode(3)
    root.left.right = TreeNode(4)
    root.right.left = TreeNode(4)
    root.right.right = TreeNode(3)

    print SolutionRecur().isSymmetric(root)
    print SolutionIter().isSymmetric(root)

    # Output: False.
    #     1
    #    / \
    #   2   2
    #    \   \
    #     3   3
    root = TreeNode(1)
    root.left = TreeNode(2)
    root.right = TreeNode(2)
    root.left.right = TreeNode(3)
    root.right.right = TreeNode(3)

    print SolutionRecur().isSymmetric(root)
    print SolutionIter().isSymmetric(root)

if __name__ == '__main__':
    main()
<commit_msg>Revise var to is_in/out/pair_mirror<commit_after>"""Leetcode 101. Symmetric Tree
Easy

URL: https://leetcode.com/problems/symmetric-tree/

Given a binary tree, check whether it is a mirror of itself
(ie, symmetric around its center).

For example, this binary tree [1,2,2,3,4,4,3] is symmetric:
    1
   / \
  2   2
 / \ / \
3  4 4  3

But the following [1,2,2,null,3,null,3] is not:
    1
   / \
  2   2
   \   \
   3    3

Note:
Bonus points if you could solve it both recursively and iteratively.
"""

# Definition for a binary tree node.
class TreeNode(object):
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None

class SolutionRecur(object):
    def isMirror(self, left, right):
        # Check left & right nodes's symmetry.
        if not left and not right:
            return True

        if not left or not right:
            return False

        # Check left & right nodes's values.
        if left.val != right.val:
            return False

        # Check if outside & inside pairs are mirrors.
        is_out_pair_mirror = self.isMirror(left.left, right.right)
        is_in_pair_mirror = self.isMirror(left.right, right.left)
        return is_out_pair_mirror and is_in_pair_mirror

    def isSymmetric(self, root):
        """
        :type root: TreeNode
        :rtype: bool

        Time complexity: O(n).
        Space complexity: O(logn) for balanced tree; O(n) for single sided tree.
        """
        if not root:
            return True

        return self.isMirror(root.left, root.right)

class SolutionIter(object):
    def isSymmetric(self, root):
        """
        :type root: TreeNode
        :rtype: bool

        Time complexity: O(n).
        Space complexity: O(logn) for balanced tree; O(n) for single sided tree.
        """
        if not root:
            return True

        queue = [(root.left, root.right)]

        while queue:
            left, right = queue.pop()

            # Check left & right nodes's symmetry.
            if not left and not right:
                continue

            if not left or not right:
                return False

            # Check left & right nodes's values.
            if left.val != right.val:
                return False

            # Check outside & inside pairs.
            queue.insert(0, (left.left, right.right))
            queue.insert(0, (left.right, right.left))

        return True

def main():
    # Output: True.
    #     1
    #    / \
    #   2   2
    #  / \ / \
    # 3  4 4  3
    root = TreeNode(1)
    root.left = TreeNode(2)
    root.right = TreeNode(2)
    root.left.left = TreeNode(3)
    root.left.right = TreeNode(4)
    root.right.left = TreeNode(4)
    root.right.right = TreeNode(3)

    print SolutionRecur().isSymmetric(root)
    print SolutionIter().isSymmetric(root)

    # Output: False.
    #     1
    #    / \
    #   2   2
    #    \   \
    #     3   3
    root = TreeNode(1)
    root.left = TreeNode(2)
    root.right = TreeNode(2)
    root.left.right = TreeNode(3)
    root.right.right = TreeNode(3)

    print SolutionRecur().isSymmetric(root)
    print SolutionIter().isSymmetric(root)

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


<commit_before># -*- coding: iso-8859-1 -*-
# -----------------------------------------------------------------------------
# main.py - Main loop functions
# -----------------------------------------------------------------------------
# $Id$
#
# -----------------------------------------------------------------------------
# kaa.notifier - Mainloop and callbacks
# Copyright (C) 2005, 2006 Dirk Meyer, Jason Tackaberry, et al.
#
# First Version: Dirk Meyer <dmeyer@tzi.de>
# Maintainer:    Dirk Meyer <dmeyer@tzi.de>
#                Jason Tackaberry <tack@urandom.ca>
#
# Please see the file AUTHORS for a complete list of authors.
#
# This library is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version
# 2.1 as published by the Free Software Foundation.
#
# This library is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301 USA
#
# -----------------------------------------------------------------------------

# python imports
import sys
import logging
import os
import time
import signal
import threading
import atexit

import nf_wrapper as notifier
from callback import Signal
from popen import proclist as _proclist
from thread import MainThreadCallback, is_mainthread, wakeup, set_mainthread
from jobserver import killall as kill_jobserver
from decorators import execute_in_mainloop

__all__ = [ 'init', 'start', 'stop', 'step', 'is_running', 'wakeup', 'set_mainthread' ]

# get logging object
log = logging.getLogger('notifier')

# variable to check if the notifier is running
_running = False
# Set if currently in shutdown() (to prevent reentrancy)
_shutting_down = False

def _step_signal_changed(signal, flag):
    if flag == Signal.SIGNAL_CONNECTED and signal.count() == 1:
        notifier.dispatcher_add(signals["step"].emit)
    elif flag == Signal.SIGNAL_DISCONNECTED and signal.count() == 0:
        notifier.dispatcher_remove(signals["step"].emit)

signals = {
    'shutdown': Signal(),
    'step': Signal(changed_cb = _step_signal_changed),
}

# Ensure stop() is called from main thread.
@execute_in_mainloop()
def stop():
    """
    Shutdown notifier and kill all background processes.
    """
    global _shutting_down

    if _running:
        # notifier loop still running, send system exit
        log.info('Stop notifier loop')
        raise SystemExit

    if _shutting_down:
        return

    _shutting_down = True

    _proclist.stop_all()
    signals["shutdown"].emit()
    signals["shutdown"].disconnect_all()
    signals["step"].disconnect_all()

    # Kill processes _after_ shutdown emits to give callbacks a chance to
    # close them properly.
    _proclist.kill_all()
    while _proclist.check():
        # wait until all processes are stopped
        step()
    kill_jobserver()
    # Collect any zombies
    try:
        os.waitpid(-1, os.WNOHANG)
    except:
        pass

def start():
    """
    Notifier main loop function. It will loop until an exception
    is raised or sys.exit is called.
    """
    global _running
    _running = True

    set_mainthread()
    try:
        while True:
            notifier.step()
    except (KeyboardInterrupt, SystemExit):
        try:
            # This looks stupid, I know that. The problem is that if we have
            # a KeyboardInterrupt, that flag is still valid somewhere inside
            # python. The next system call will fail because of that. Since we
            # don't want a join of threads or similar fail, we use a very short
            # sleep here. In most cases we won't sleep at all because this sleep
            # fails. But after that everything is back to normal.
            time.sleep(0.001)
        except:
            pass
    except Exception, e:
        log.exception('loop')
    _running = False
    stop()

def is_running():
    return _running

def _set_running(status):
    global _running
    _running = status

def select_notifier(module, **options):
    """
    Initialize the specified notifier.
    """
    if module in ('thread', 'twisted'):
        import nf_thread
        return nf_thread.init(module, **options)
    return notifier.init( module, **options )

def step(*args, **kwargs):
    """
    Notifier step function with signal support.
    """
    if not is_mainthread():
        # If step is being called from a thread, wake up the mainthread
        # instead of allowing the thread into notifier.step.
        wakeup()
        # Sleep for epsilon to prevent busy loops.
        time.sleep(0.001)
        return

    try:
        notifier.step(*args, **kwargs)
    except (KeyboardInterrupt, SystemExit):
        raise SystemExit

def _shutdown_check(*args):
    # Helper function to shutdown kaa on system exit
    # The problem is that pytgtk just exits python and
    # does not simply return from the main loop and kaa
    # can't call the shutdown handler. This is not a perfect
    # solution, e.g. with the generic notifier you can do
    # stuff after kaa.main() which is not possible with gtk
    global _running
    if _running:
        # If the kaa mainthread (i.e. thread the mainloop is running in)
        # is not the program's main thread, then is_mainthread() will be False
        # and we don't need to set running=False since shutdown() will raise a
        # SystemExit and things will exit normally.
        if is_mainthread():
            _running = False
        stop()

# catch SIGTERM and SIGINT if possible for a clean shutdown
if threading.enumerate()[0] == threading.currentThread():
    def signal_handler(*args):
        sys.exit(0)
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)
else:
    log.info('kaa imported from thread, disable SIGTERM handler')

# check to make sure we really call our shutdown function
atexit.register(_shutdown_check)
<commit_msg>Don't wait for result on shutdown() if not called from mainthread.<commit_after># -*- coding: iso-8859-1 -*-
# -----------------------------------------------------------------------------
# main.py - Main loop functions
# -----------------------------------------------------------------------------
# $Id$
#
# -----------------------------------------------------------------------------
# kaa.notifier - Mainloop and callbacks
# Copyright (C) 2005, 2006 Dirk Meyer, Jason Tackaberry, et al.
#
# First Version: Dirk Meyer <dmeyer@tzi.de>
# Maintainer:    Dirk Meyer <dmeyer@tzi.de>
#                Jason Tackaberry <tack@urandom.ca>
#
# Please see the file AUTHORS for a complete list of authors.
#
# This library is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version
# 2.1 as published by the Free Software Foundation.
#
# This library is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301 USA
#
# -----------------------------------------------------------------------------

# python imports
import sys
import logging
import os
import time
import signal
import threading
import atexit

import nf_wrapper as notifier
from callback import Signal
from popen import proclist as _proclist
from thread import MainThreadCallback, is_mainthread, wakeup, set_mainthread
from jobserver import killall as kill_jobserver
from decorators import execute_in_mainloop

__all__ = [ 'init', 'start', 'stop', 'step', 'is_running', 'wakeup', 'set_mainthread' ]

# get logging object
log = logging.getLogger('notifier')

# variable to check if the notifier is running
_running = False
# Set if currently in shutdown() (to prevent reentrancy)
_shutting_down = False

def _step_signal_changed(signal, flag):
    if flag == Signal.SIGNAL_CONNECTED and signal.count() == 1:
        notifier.dispatcher_add(signals["step"].emit)
    elif flag == Signal.SIGNAL_DISCONNECTED and signal.count() == 0:
        notifier.dispatcher_remove(signals["step"].emit)

signals = {
    'shutdown': Signal(),
    'step': Signal(changed_cb = _step_signal_changed),
}

# Ensure stop() is called from main thread.
@execute_in_mainloop(async = True)
def stop():
    """
    Shutdown notifier and kill all background processes.
    """
    global _shutting_down

    if _running:
        # notifier loop still running, send system exit
        log.info('Stop notifier loop')
        raise SystemExit

    if _shutting_down:
        return

    _shutting_down = True

    _proclist.stop_all()
    signals["shutdown"].emit()
    signals["shutdown"].disconnect_all()
    signals["step"].disconnect_all()

    # Kill processes _after_ shutdown emits to give callbacks a chance to
    # close them properly.
    _proclist.kill_all()
    while _proclist.check():
        # wait until all processes are stopped
        step()
    kill_jobserver()
    # Collect any zombies
    try:
        os.waitpid(-1, os.WNOHANG)
    except:
        pass

def start():
    """
    Notifier main loop function. It will loop until an exception
    is raised or sys.exit is called.
    """
    global _running
    _running = True

    set_mainthread()
    try:
        while True:
            notifier.step()
    except (KeyboardInterrupt, SystemExit):
        try:
            # This looks stupid, I know that. The problem is that if we have
            # a KeyboardInterrupt, that flag is still valid somewhere inside
            # python. The next system call will fail because of that. Since we
            # don't want a join of threads or similar fail, we use a very short
            # sleep here. In most cases we won't sleep at all because this sleep
            # fails. But after that everything is back to normal.
            time.sleep(0.001)
        except:
            pass
    except Exception, e:
        log.exception('loop')
    _running = False
    stop()

def is_running():
    return _running

def _set_running(status):
    global _running
    _running = status

def select_notifier(module, **options):
    """
    Initialize the specified notifier.
    """
    if module in ('thread', 'twisted'):
        import nf_thread
        return nf_thread.init(module, **options)
    return notifier.init( module, **options )

def step(*args, **kwargs):
    """
    Notifier step function with signal support.
    """
    if not is_mainthread():
        # If step is being called from a thread, wake up the mainthread
        # instead of allowing the thread into notifier.step.
        wakeup()
        # Sleep for epsilon to prevent busy loops.
        time.sleep(0.001)
        return

    try:
        notifier.step(*args, **kwargs)
    except (KeyboardInterrupt, SystemExit):
        raise SystemExit

def _shutdown_check(*args):
    # Helper function to shutdown kaa on system exit
    # The problem is that pytgtk just exits python and
    # does not simply return from the main loop and kaa
    # can't call the shutdown handler. This is not a perfect
    # solution, e.g. with the generic notifier you can do
    # stuff after kaa.main() which is not possible with gtk
    global _running
    if _running:
        # If the kaa mainthread (i.e. thread the mainloop is running in)
        # is not the program's main thread, then is_mainthread() will be False
        # and we don't need to set running=False since shutdown() will raise a
        # SystemExit and things will exit normally.
        if is_mainthread():
            _running = False
        stop()

# catch SIGTERM and SIGINT if possible for a clean shutdown
if threading.enumerate()[0] == threading.currentThread():
    def signal_handler(*args):
        sys.exit(0)
    signal.signal(signal.SIGTERM, signal_handler)
    signal.signal(signal.SIGINT, signal_handler)
else:
    log.info('kaa imported from thread, disable SIGTERM handler')

# check to make sure we really call our shutdown function
atexit.register(_shutdown_check)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import sys, os, codecs, random
from datetime import datetime

import yaml
import PyRSS2Gen.PyRSS2Gen as RSS2Gen

import util
from setting import *

"""
文章信息

layout   : 布局名称
title    : 标题
date     : 发布时间
tags     : 标签
category : 分类

source   : 原始文件名
content  : 内容
"""

""" 单篇文章 """
class Article(object):

    def __init__(self, site, source):
        self.site = site
        self.source = source

        self.layout = 'post'
        self.title = 'untitled'
        self.date = None
        self.author = self.site.siteAuthor
        self.tags = []
        self.category = []

        self.content = ''
        self.summary = ''

        self.metadata = {}

        self.unique = ''.join(random.sample('9876543210zyxwvutsrqponmlkjihgfedcba', 5))

    def parse(self):
        if not os.path.exists(self.source):
            util.log('file non-existent. [%s]' % self.source)
            return False

        fp = codecs.open(self.source, 'r', encoding='utf-8')
        lines = fp.readlines()

        # 解析文章信息
        config_str = ''
        config = {}
        index = 1
        for line in lines[1:]:
            index += 1
            if line.find('---') == 0:
                break
            config_str += line
        config.update( yaml.load(config_str) )

        self.layout     = config.get('layout', self.site.defaultLayout)
        self.title      = config.get('title', 'untitled')
        self.author     = config.get('author', self.site.siteAuthor)
        self.date       = config.get('date', None)
        self.tag        = config.get('tag', [])
        self.category   = config.get('category', [])
        self.metadata   = config.get('metadata', {})

        # 检查date
        if not self.checkDateAvailable():
            print "invalid article: date error '%s'" % self.site.getRelativePath(self.source)
            return False

        # 检查分类
        self.checkCategoryAvailable()

        # 获取内容
        self.content = ''
        self.summary = ''
        for line in lines[index:]:
            if line.strip(' \t').find('') == 0:
                self.summary = self.content
                line = '' % self.unique
            self.content += line

        self.content.strip('\n\t ')
        self.summary.strip('\n\t ')

        if len(self.content) == 0:  #? 都不会成立，总是会有一个类似换行的东西 又不是\n WHY?
            print "invalid article: content empty '%s'" % self.site.getRelativePath(self.source)
            return False

        # 处理 html 与 markdown 格式
        if self.isMarkdown():
            self.content = util.markdown(self.content)
            self.summary = util.markdown(self.summary)

        # 摘要
        if not self.summary:
            self.summary = self.content

        return True

    # 文章日期是否有效
    def checkDateAvailable(self):
        if isinstance(self.date, datetime):
            return True
        if not self.date or not isinstance(self.date, str):
            return False
        date = None
        for p in ARTICLE_DATE_FORMAT:
            try:
                date = datetime.strptime(self.date, p)
            except Exception, e:
                continue
            else:
                self.date = date
                break
        if not isinstance(self.date, datetime):
            return False
        if self.date > datetime.now():
            return False
        return True

    # 检查文章分类的有效性 只有在网站配置'category'中存在了才可用
    def checkCategoryAvailable(self):
        cate = []
        for c in self.category:
            category = self.site.convertToConfigedCategory(c)
            if category:
                if category not in cate:
                    cate.append(category)
                else:
                    print "category: '%s' already exists IN '%s'" % (c, self.site.getRelativePath(self.source))
            else:
                print "unavailable category: '%s' IN %s" % (c, self.site.getRelativePath(self.source))
        self.category = cate

    def export(self):
        print '  %s' % self.site.getRelativePath(self.source)
        data = {'site': self.site, 'article': self}
        self.site.exportFile(self.exportFilePath, self.layout, data)

    def isMarkdown(self):
        extension = os.path.splitext(self.source)[1][1:]
        return extension.find('md') == 0 or extension.find('markdown') == 0

    # 永久链接
    @property
    def permalink(self):
        return self.site.generateUrl(self.exportDir, self.exportBasename)

    # 跳转到 more 的链接
    @property
    def morePermalink(self):
        return '%s#%s' % (self.permalink, self.unique)

    # 导出文件的绝对路径
    @property
    def exportFilePath(self):
        return self.site.generateDeployFilePath(self.exportDir, self.exportBasename)

    @property
    def exportDir(self):
        #! 当 self.source 包括子文件夹时使用自身 如果没有则统一放入article目录
        dirname, basename = os.path.split(self.source)
        dirname = dirname[len( self.site.articlePath )+1:].strip(' /')
        if not dirname:
            dirname += 'article'
        return dirname

    # 源文件的文件名（不包括后缀）
    @property
    def exportBasename(self):
        source_basename = os.path.splitext(os.path.split(self.source)[1])[0]
        if self.site.isArticleFilenameUseDatePrefix:
            return source_basename[ARTICLE_FILENAME_PREFIX_LEN:]
        return source_basename

    def getFeedItem(self):
        return RSS2Gen.RSSItem(
            title = self.title,
            link = self.permalink,
            guid = RSS2Gen.Guid(self.permalink),
            description = self.summary,
            pubDate = self.date,
            #categories = tags,
            )

""" 页面导航信息 """
class PageNav(object):
    def __init__(self, articles, num_per_page, first_page_name, other_page_name = 'page-%d'):
        assert isinstance(articles, list) and len(articles) > 0 and num_per_page > 0
        self.articles = articles
        self.num_per_page = num_per_page
        self.first_page_name = first_page_name
        self.other_page_name = other_page_name

    # 页面数
    def getPageCount(self):
        count = len(self.articles) / self.num_per_page;
        if len(self.articles) % self.num_per_page > 0:
            count += 1
        return count

    # 获取特定页面的数据 [articles, deploy_file_name]
    def getDataByPageNum(self, num):
        assert 0 < num <= self.getPageCount()
        start_index = self.num_per_page * (num-1)
        end_index = start_index + self.num_per_page + 1
        deploy_file_name = self.first_page_name;
        if num > 1:
            deploy_file_name = self.other_page_name % num
        return [self.articles[start_index:end_index], deploy_file_name]

class ArticleBundle(object):
    """文章集合
    有多篇文章集合的处理  存档 首页 分类 标签
    """
    def __init__(self, site):
        self.site = site
        self.articles = []

    def addArticle(self, article):
        if not isinstance (article, Article):
            raise ValueError, 'type of article is not Article'
        elif article in self.articles:
            raise ValueError, 'article is already exists'
        self.articles.append(article)

    def export(self):
        self.printSelf()

        if self.count <= 0:
            return

        self.sortByDateDESC()

        paged_articles = self.getPagedArticles()
        total_page_num = self.totalPageNum

        for i in range(1, total_page_num + 1):

            articles = paged_articles[i - 1]
            deploy_file = self.site.generateDeployFilePath(self.exportDir, page=i)

            self.curPageNum = i
            data = {'site': self.site, 'articles': articles}
            self.site.exportFile(deploy_file, self.templateName, data)

    def printSelf(self):
        return

    # 按日期倒序排序
    def sortByDateDESC(self):
        self.articles = sorted(self.articles, key = lambda a: a.date, reverse = True )
        pass

    @property
    def count(self):
        return len(self.articles)

    @property
    def permalink(self):
        return self.site.generateUrl(self.exportDir)

    def getPagePermalink(self, page_num):
        assert 0 < page_num <= self.totalPageNum, 'page number must in 1 - %d' % self.totalPageNum
        return self.site.generateUrl(self.exportDir, page_num)

    @property
    def templateName(self):
        return ''

    # 输出目录 相对deploy
    @property
    def exportDir(self):
        return ''

    # 每页文章数目
    @property
    def numPerPage(self):
        return self.site.numPerPage

    @property
    def totalPageNum(self):
        count = len(self.articles) / self.numPerPage;
        if len(self.articles) % self.numPerPage > 0:
            count += 1
        return count

    def getPagedArticles(self):
        paged = []
        total = self.totalPageNum
        for i in range(0, total):
            start = self.numPerPage * i
            if i == total - 1:
                paged.append( self.articles[start:] )
            else:
                paged.append( self.articles[start:self.numPerPage] )
        return paged

class Archive(ArticleBundle):
    """ 存档 存档页 文章单页的输出 """
    def __init__(self, site, articles):
        super(Archive, self).__init__(site)
        self.articles = articles

    def printSelf(self):
        print 'Archive:'

    @property
    def templateName(self):
        return 'archive'

    @property
    def exportDir(self):
        return 'archive'

    @property
    def numPerPage(self):
        return self.site.numPerPageInArchive;

class Home(ArticleBundle):
    """ 首页的输出 """
    def __init__(self, site, articles):
        super(Home, self).__init__(site)
        self.articles = articles

    def printSelf(self):
        print 'Home: %s' % self.permalink

    @property
    def templateName(self):
        return 'home'

class Category(ArticleBundle):
    """ 分类 """
    def __init__(self, site, category):
        super(Category, self).__init__(site)
        self.category_name = category

    def printSelf(self):
        print 'Category: %s %d %s' % (self.category_name, self.count, self.permalink)

    @property
    def templateName(self):
        return 'category'

    @property
    def exportDir(self):
        return 'category/%s' % self.category_name

class Tag(ArticleBundle):
    """标签
    标签页的输出
    """
    def __init__(self, site, tag):
        super(Tag, self).__init__(site)
        self.tag_name = tag

    def printSelf(self):
        print 'Tag: %s %d %s' % (self.tag_name, self.count, self.permalink)

    @property
    def templateName(self):
        return 'tag'

    @property
    def exportDir(self):
        return 'tag/%s' % self.tag_name

class Feed(ArticleBundle):
    """ Feed的输出 """
    def __init__(self, site, articles):
        super(Feed, self).__init__(site)
        self.articles = articles

    # 忽略输出方法
    def export(self):

        self.sortByDateDESC()
        num = self.site.numInFeed;
        articles = []
        if len(self.articles) >= num:
            articles = self.articles[0:num]
        else:
            articles = self.articles[0:]

        feed = RSS2Gen.RSS2(
            title = self.site.siteTitle,
            link = self.site.siteUrl,
            description = self.site.siteTagline,
            lastBuildDate = datetime.now(),
            items = [ a.getFeedItem() for a in articles ],
            generator = util.self()
            )

        filename = self.site.generateDeployFilePath(self.site.feedFilename, assign=True)
        with open(filename, 'w') as fp:
                feed.write_xml(fp, 'utf-8')

class Tags(ArticleBundle):
    """标签"""
    def __init__(self, site, articles):
        self.site = site
        self.articles = articles
        self.tags = self.fetchAllTags()

    def fetchAllTags(self):
        tags = {}
        for article in self.articles:
            for tag in article.tag:
                slug = util.generateSlug(tag)
                if slug in tags.keys():
                    tags[slug].addArticle(article)
                else:
                    tags[slug] = Tag(self.site, tag)
                    tags[slug].addArticle(article)
        # 按文章数排序
        return sorted(tags.values(), cmp = lambda a, b: a.count > b.count)

    def export(self):
        self.printSelf()

        if self.count == 0:
            return

        data = {'site': self.site, 'tags': self.tags}
        deploy_file = self.site.generateDeployFilePath('tags')

        self.site.exportFile(deploy_file, self.templateName, data)

        # 输出标签单页
        map(lambda t: t.export(), self.tags)

    def printSelf(self):
        print 'Tags: %s' % self.permalink

    @property
    def count(self):
        return len(self.tags)

    @property
    def exportDir(self):
        return 'tags'

    @property
    def templateName(self):
        return 'tags'

class Categories(ArticleBundle):
    """分类"""
    def __init__(self, site, articles):
        super(Categories, self).__init__(site)
        self.articles = articles
        self.categories = self.fetchAllCategories()

    def fetchAllCategories(self):
        categories = {}
        for article in self.articles:
            for category in article.category:
                # 分类是否满足条件在 文章对象中处理
                slug = util.generateSlug(category)
                if slug in categories.keys():
                    categories[slug].addArticle(article)
                else:
                    categories[slug] = Category(self.site, category)
                    categories[slug].addArticle(article)
        # 按文章数排序
        return sorted(categories.values(), cmp = lambda a, b: a.count > b.count)

    def export(self):
        self.printSelf()
        if not self.count:
            return
        map(lambda c: c.export(), self.categories)

    def printSelf(self):
        print 'Categories: %s' % self.permalink

    @property
    def count(self):
        return len(self.categories)

        <commit_msg>Tag Category 使用对象存储<commit_after># -*- coding: utf-8 -*-
import sys, os, codecs, random
from datetime import datetime

import yaml
import PyRSS2Gen.PyRSS2Gen as RSS2Gen

import util
from setting import *

"""
文章信息

layout   : 布局名称
title    : 标题
date     : 发布时间
tags     : 标签
category : 分类

source   : 原始文件名
content  : 内容
"""

""" 单篇文章 """
class Article(object):

    def __init__(self, site, source):
        self.site = site
        self.source = source

        self.layout = 'post'
        self.title = 'untitled'
        self.date = None
        self.author = self.site.siteAuthor
        self.tag = []
        self.category = []

        self.content = ''
        self.summary = ''

        self.metadata = {}

        self.unique = ''.join(random.sample('9876543210zyxwvutsrqponmlkjihgfedcba', 5))

    def parse(self):
        if not os.path.exists(self.source):
            util.log('file non-existent. [%s]' % self.source)
            return False

        fp = codecs.open(self.source, 'r', encoding='utf-8')
        lines = fp.readlines()

        # 解析文章信息
        config_str = ''
        config = {}
        index = 1
        for line in lines[1:]:
            index += 1
            if line.find('---') == 0:
                break
            config_str += line
        config.update( yaml.load(config_str) )

        self.layout     = config.get('layout', self.site.defaultLayout)
        self.title      = config.get('title', 'untitled')
        self.author     = config.get('author', self.site.siteAuthor)
        self.date       = config.get('date', None)
        self.tag        = config.get('tag', [])
        self.category   = config.get('category', [])
        self.metadata   = config.get('metadata', {})

        # 检查date
        if not self.checkDateAvailable():
            print "invalid article: date error '%s'" % self.site.getRelativePath(self.source)
            return False

        # 解析分类
        self.parseCategory()
        self.parseTag()

        # 获取内容
        self.content = ''
        self.summary = ''
        for line in lines[index:]:
            if line.strip(' \t').find('') == 0:
                self.summary = self.content
                line = '' % self.unique
            self.content += line

        self.content.strip('\n\t ')
        self.summary.strip('\n\t ')

        if len(self.content) == 0:  #? 都不会成立，总是会有一个类似换行的东西 又不是\n WHY?
            print "invalid article: content empty '%s'" % self.site.getRelativePath(self.source)
            return False

        # 处理 html 与 markdown 格式
        if self.isMarkdown():
            self.content = util.markdown(self.content)
            self.summary = util.markdown(self.summary)

        # 摘要
        if not self.summary:
            self.summary = self.content

        return True

    # 文章日期是否有效
    def checkDateAvailable(self):
        if isinstance(self.date, datetime):
            return True
        if not self.date or not isinstance(self.date, str):
            return False
        date = None
        for p in ARTICLE_DATE_FORMAT:
            try:
                date = datetime.strptime(self.date, p)
            except Exception, e:
                continue
            else:
                self.date = date
                break
        if not isinstance(self.date, datetime):
            return False
        if self.date > datetime.now():
            return False
        return True

    # 检查文章分类的有效性 只有在网站配置'category'中存在了才可用
    # 对象转换为Category对象
    def parseCategory(self):
        cate = []
        for c in self.category:
            category = self.site.convertToConfigedCategory(c)
            if category:
                if category not in cate:
                    cate_obj = Category(self.site, category)
                    cate.append(cate_obj)
                else:
                    print "category: '%s' already exists IN '%s'" % (c, self.site.getRelativePath(self.source))
            else:
                print "unavailable category: '%s' IN %s" % (c, self.site.getRelativePath(self.source))
        self.category = cate

    def parseTag(self):
        tag = []
        for t in self.tag:
            if t:
                tag_obj = Tag(self.site, t)
                tag.append(tag_obj)
        self.tag = tag

    def export(self):
        print '  %s' % self.site.getRelativePath(self.source)
        data = {'site': self.site, 'article': self}
        self.site.exportFile(self.exportFilePath, self.layout, data)

    def isMarkdown(self):
        extension = os.path.splitext(self.source)[1][1:]
        return extension.find('md') == 0 or extension.find('markdown') == 0

    # 永久链接
    @property
    def permalink(self):
        return self.site.generateUrl(self.exportDir, self.exportBasename)

    # 跳转到 more 的链接
    @property
    def morePermalink(self):
        return '%s#more-%s' % (self.permalink, self.unique)

    # 导出文件的绝对路径
    @property
    def exportFilePath(self):
        return self.site.generateDeployFilePath(self.exportDir, self.exportBasename)

    @property
    def exportDir(self):
        #! 当 self.source 包括子文件夹时使用自身 如果没有则统一放入article目录
        dirname, basename = os.path.split(self.source)
        dirname = dirname[len( self.site.articlePath )+1:].strip(' /')
        if not dirname:
            dirname += 'article'
        return dirname

    # 源文件的文件名（不包括后缀）
    @property
    def exportBasename(self):
        source_basename = os.path.splitext(os.path.split(self.source)[1])[0]
        if self.site.isArticleFilenameUseDatePrefix:
            return source_basename[ARTICLE_FILENAME_PREFIX_LEN:]
        return source_basename

    def getFeedItem(self):
        return RSS2Gen.RSSItem(
            title = self.title,
            link = self.permalink,
            guid = RSS2Gen.Guid(self.permalink),
            description = self.summary,
            pubDate = self.date,
            #categories = tags,
            )

    def outputSummary(self, text='Read more...', format='
%s
', css='more-link'):
        if self.summary == self.content:
            return self.content
        more_link = '%s' % (self.morePermalink, css, text)
        return self.summary + '\n' + format % more_link

""" 页面导航信息 """
class PageNav(object):
    def __init__(self, articles, num_per_page, first_page_name, other_page_name = 'page-%d'):
        assert isinstance(articles, list) and len(articles) > 0 and num_per_page > 0
        self.articles = articles
        self.num_per_page = num_per_page
        self.first_page_name = first_page_name
        self.other_page_name = other_page_name

    # 页面数
    def getPageCount(self):
        count = len(self.articles) / self.num_per_page;
        if len(self.articles) % self.num_per_page > 0:
            count += 1
        return count

    # 获取特定页面的数据 [articles, deploy_file_name]
    def getDataByPageNum(self, num):
        assert 0 < num <= self.getPageCount()
        start_index = self.num_per_page * (num-1)
        end_index = start_index + self.num_per_page + 1
        deploy_file_name = self.first_page_name;
        if num > 1:
            deploy_file_name = self.other_page_name % num
        return [self.articles[start_index:end_index], deploy_file_name]

class ArticleBundle(object):
    """文章集合
    有多篇文章集合的处理  存档 首页 分类 标签
    """
    def __init__(self, site):
        self.site = site
        self.articles = []

    def addArticle(self, article):
        if not isinstance (article, Article):
            raise ValueError, 'type of article is not Article'
        elif article in self.articles:
            raise ValueError, 'article is already exists'
        self.articles.append(article)

    def export(self):
        self.printSelf()

        if self.count <= 0:
            return

        self.sortByDateDESC()

        paged_articles = self.getPagedArticles()
        total_page_num = self.totalPageNum

        for i in range(1, total_page_num + 1):
            articles = paged_articles[i - 1]
            deploy_file = self.site.generateDeployFilePath(self.exportDir, page=i)
            self.curPageNum = i
            data = {'site': self.site, 'articles': articles}
            self.site.exportFile(deploy_file, self.templateName, data)

    def printSelf(self):
        return

    # 按日期倒序排序
    def sortByDateDESC(self):
        self.articles = sorted(self.articles, key = lambda a: a.date, reverse = True )
        pass

    @property
    def count(self):
        return len(self.articles)

    @property
    def permalink(self):
        return self.site.generateUrl(self.exportDir)

    def getPagePermalink(self, page_num):
        assert 0 < page_num <= self.totalPageNum, 'page number must in 1 - %d' % self.totalPageNum
        return self.site.generateUrl(self.exportDir, page_num)

    @property
    def templateName(self):
        return ''

    # 输出目录 相对deploy
    @property
    def exportDir(self):
        return ''

    # 每页文章数目
    @property
    def numPerPage(self):
        return self.site.numPerPage

    @property
    def totalPageNum(self):
        count = len(self.articles) / self.numPerPage;
        if len(self.articles) % self.numPerPage > 0:
            count += 1
        return count

    def getPagedArticles(self):
        paged = []
        total = self.totalPageNum
        for i in range(0, total):
            start = self.numPerPage * i
            end = start + self.numPerPage
            if i == total - 1:
                paged.append( self.articles[start:] )
            else:
                paged.append( self.articles[start:end] )
        return paged

class Archive(ArticleBundle):
    """ 存档 存档页 文章单页的输出 """
    def __init__(self, site, articles):
        super(Archive, self).__init__(site)
        self.articles = articles

    def printSelf(self):
        print 'Archive:'

    @property
    def templateName(self):
        return 'archive'

    @property
    def exportDir(self):
        return 'archive'

    @property
    def numPerPage(self):
        return self.site.numPerPageInArchive;

class Home(ArticleBundle):
    """ 首页的输出 """
    def __init__(self, site, articles):
        super(Home, self).__init__(site)
        self.articles = articles

    def printSelf(self):
        print 'Home: %s' % self.permalink

    @property
    def templateName(self):
        return 'home'

class Category(ArticleBundle):
    """ 分类 """
    def __init__(self, site, category):
        super(Category, self).__init__(site)
        self.name = category

    def printSelf(self):
        print 'Category: %s %d %s' % (self.name, self.count, self.permalink)

    @property
    def templateName(self):
        return 'category'

    @property
    def exportDir(self):
        return 'category/%s' % self.name

class Tag(ArticleBundle):
    """标签
    标签页的输出
    """
    def __init__(self, site, tag):
        super(Tag, self).__init__(site)
        self.name = tag

    def printSelf(self):
        print 'Tag: %s %d %s' % (self.name, self.count, self.permalink)

    @property
    def templateName(self):
        return 'tag'

    @property
    def exportDir(self):
        return 'tag/%s' % self.name

class Feed(ArticleBundle):
    """ Feed的输出 """
    def __init__(self, site, articles):
        super(Feed, self).__init__(site)
        self.articles = articles

    # 忽略输出方法
    def export(self):

        self.sortByDateDESC()
        num = self.site.numInFeed;
        articles = []
        if len(self.articles) >= num:
            articles = self.articles[0:num]
        else:
            articles = self.articles[0:]

        feed = RSS2Gen.RSS2(
            title = self.site.siteTitle,
            link = self.site.siteUrl,
            description = self.site.siteTagline,
            lastBuildDate = datetime.now(),
            items = [ a.getFeedItem() for a in articles ],
            generator = util.self()
            )

        filename = self.site.generateDeployFilePath(self.site.feedFilename, assign=True)
        with open(filename, 'w') as fp:
                feed.write_xml(fp, 'utf-8')

class Tags(ArticleBundle):
    """标签"""
    def __init__(self, site, articles):
        self.site = site
        self.articles = articles
        self.tags = self.fetchAllTags()

    def fetchAllTags(self):
        tags = {}
        for article in self.articles:
            for tag in article.tag:
                slug = util.generateSlug(tag.name)
                if slug in tags.keys():
                    tags[slug].addArticle(article)
                else:
                    tags[slug] = Tag(self.site, tag.name)
                    tags[slug].addArticle(article)
        # 按文章数排序
        return sorted(tags.values(), cmp = lambda a, b: a.count > b.count)

    def export(self):
        self.printSelf()

        if self.count == 0:
            return

        data = {'site': self.site, 'tags': self.tags}
        deploy_file = self.site.generateDeployFilePath('tags')

        self.site.exportFile(deploy_file, self.templateName, data)

        # 输出标签单页
        map(lambda t: t.export(), self.tags)

    def printSelf(self):
        print 'Tags: %s' % self.permalink

    @property
    def count(self):
        return len(self.tags)

    @property
    def exportDir(self):
        return 'tags'

    @property
    def templateName(self):
        return 'tags'

class Categories(ArticleBundle):
    """分类"""
    def __init__(self, site, articles):
        super(Categories, self).__init__(site)
        self.articles = articles
        self.categories = self.fetchAllCategories()

    def fetchAllCategories(self):
        categories = {}
        for article in self.articles:
            for category in article.category:
                category_name = category.name
                # 分类是否满足条件在 文章对象中处理
                slug = util.generateSlug(category_name)
                if slug in categories.keys():
                    categories[slug].addArticle(article)
                else:
                    categories[slug] = Category(self.site, category_name)
                    categories[slug].addArticle(article)
        # 按文章数排序
        return sorted(categories.values(), cmp = lambda a, b: a.count > b.count)

    def export(self):
        self.printSelf()
        if not self.count:
            return
        map(lambda c: c.export(), self.categories)

    def printSelf(self):
        print 'Categories: %s' % self.permalink

    @property
    def count(self):
        return len(self.categories)

        <|endoftext|>


<commit_before>import argparse
import json
import os.path
import re
from time import strftime, gmtime

class PrettyOuput(object):

    def d2b(self, a):
        tbin = []
        while a:
            tbin.append(a % 2)
            a //= 2

        t2bin = tbin[::-1]
        if len(t2bin) != 8:
            for x in range(6 - len(t2bin)):
                t2bin.insert(0, 0)
        return ''.join([str(g) for g in t2bin])

    def convert(self, time):
        if time == 0:
            return "None"
        if time == -9223372036854775808:
            return "Not Set"
        sec = abs(time) // 10000000
        days = sec // 86400
        sec -= 86400*days
        hrs = sec // 3600
        sec -= 3600*hrs
        mins = sec // 60
        sec -= 60*mins
        result = ""
        if days > 1:
            result += "{0} days ".format(days)
        elif days == 1:
            result += "{0} day ".format(days)
        if hrs > 1:
            result += "{0} hours ".format(hrs)
        elif hrs == 1:
            result += "{0} hour ".format(hrs)
        if mins > 1:
            result += "{0} minutes ".format(mins)
        elif mins == 1:
            result += "{0} minute ".format(mins)
        return result

    def password_complexity(self, data):

        print('''
    +-----------------------------------------+
    | Password Policy Information             |
    +-----------------------------------------+
    ''')

        print("[+] Password Info for Domain:", data[0]['attributes']['dc'][0].upper())
        print("\t[+] Minimum password length: ", data[0]['attributes']['instanceType'][0])
        print("\t[+] Password history length:", data[0]['attributes']['pwdHistoryLength'][0])

        password_properties = self.d2b(data[0]['attributes']['pwdProperties'][0])
        print("\t[+] Password Complexity Flags:", password_properties)
        print("")
        print("\t\t[+] Domain Refuse Password Change:", password_properties[0])
        print("\t\t[+] Domain Password Store Cleartext:", password_properties[1])
        print("\t\t[+] Domain Password Lockout Admins:", password_properties[2])
        print("\t\t[+] Domain Password No Clear Change:", password_properties[3])
        print("\t\t[+] Domain Password No Anon Change:", password_properties[4])
        print("\t\t[+] Domain Password Complex:", password_properties[5])
        print("")
        print("\t[+] Maximum password age:", self.convert(data[0]['attributes']['maxPwdAge'][0]))
        print("\t[+] Minimum password age:", self.convert(data[0]['attributes']['minPwdAge'][0]))
        print("\t[+] Reset Account Lockout Counter:", self.convert(data[0]['attributes']['lockoutDuration'][0]))
        print("\t[+] Account Lockout Threshold:", data[0]['attributes']['lockoutThreshold'][0])
        print("\t[+] Forced Log off Time:", self.convert(data[0]['attributes']['forceLogoff'][0]))

    def domain_info(self, data):
        print('''
    +--------------------------------------+
    | Getting Domain Sid For               |
    +--------------------------------------+
    ''')
        print('[+] Domain Name:', data[0]['attributes']['dc'][0])
        print('Domain Sid:', data[0]['attributes']['objectSid'][0])
        print('')
        return data[0]['attributes']['dc'][0]

    def user_info(self, users, dc):
        print('''
    +------------------------+
    | Users Infos            |
    +------------------------+
    ''')
        for user in users:
            desc = user['attributes'].get('description')[0] if user['attributes'].get('description') else "(null)"
            print("Account: " + dc + "\\" + user['attributes']['sAMAccountName'][0] + "\tName: " + user['attributes']['name'][0] + "\tDesc: " + desc)

        print("")
        for user in users:
            print("user:[" + user['attributes']['sAMAccountName'][0] + "]")
        print("")

    def groups_info(self, groups, dc):
        print('''
    +------------------------+import os.path
    | Groups Infos           |
    +------------------------+
    ''')
        for group in groups:
            print("group:[" + group['attributes']['name'][0] + "]")

        for group in groups:
            if group['attributes'].get('member'):
                users = re.findall(r"^CN=([\w\s\-\_\{\}\.\$\#]+)", '\n'.join(group['attributes']['member']), re.M)
                if users:
                    print("\n[+] Getting domain group memberships:")
                    for user in users:
                        if user == "S-1-5-11":
                            user = "NT AUTHORITY\Authenticated Users"
                        elif user == "S-1-5-4":
                            user = "NT AUTHORITY\INTERACTIVE"
                        elif user == "S-1-5-17":
                            user = "NT AUTHORITY\IUSR"
                        print("Group '" + group['attributes']['name'][0] + "' has member: " + dc + "\\" + user)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='LDAPDomainDump to pretty output like enum4linux.')

    #Main parameters
    parser.add_argument("-d", "--directory", help="The ldapdomaindump directory where the json files are saved. Required files: domain_users.json, domain_groups.json and domain_policy.json")
    args = parser.parse_args()

    if args.directory:
        with open(args.directory + '/domain_policy.json') as f:
            domain = json.load(f)
        with open(args.directory + '/domain_users.json') as f:
            users = json.load(f)
        with open(args.directory + '/domain_groups.json') as f:
            groups = json.load(f)
        pretty = PrettyOuput()
        dc = pretty.domain_info(domain)
        pretty.password_complexity(domain)
        pretty.user_info(users, dc.upper())
        pretty.groups_info(groups, dc.upper())
    else:
        print("Missing parameter --directory /output/")<commit_msg>py2 / py3 fixes<commit_after>from __future__ import print_function
from builtins import str
import argparse
import json
import os.path
import re
from time import strftime, gmtime

class PrettyOuput(object):

    def d2b(self, a):
        tbin = []
        while a:
            tbin.append(a % 2)
            a //= 2

        t2bin = tbin[::-1]
        if len(t2bin) != 8:
            for x in range(6 - len(t2bin)):
                t2bin.insert(0, 0)
        return ''.join([str(g) for g in t2bin])

    def convert(self, time):
        if isinstance(time, str):
            return time
        if time == 0:
            return "None"
        if time == -9223372036854775808:
            return "Not Set"
        sec = abs(time) // 10000000
        days = sec // 86400
        sec -= 86400*days
        hrs = sec // 3600
        sec -= 3600*hrs
        mins = sec // 60
        sec -= 60*mins
        result = ""
        if days > 1:
            result += "{0} days ".format(days)
        elif days == 1:
            result += "{0} day ".format(days)
        if hrs > 1:
            result += "{0} hours ".format(hrs)
        elif hrs == 1:
            result += "{0} hour ".format(hrs)
        if mins > 1:
            result += "{0} minutes ".format(mins)
        elif mins == 1:
            result += "{0} minute ".format(mins)
        return result

    def password_complexity(self, data):

        print('''
    +-----------------------------------------+
    | Password Policy Information             |
    +-----------------------------------------+
    ''')

        print("[+] Password Info for Domain:", data[0]['attributes']['dc'][0].upper())
        print("\t[+] Minimum password length: ", data[0]['attributes']['instanceType'][0])
        print("\t[+] Password history length:", data[0]['attributes']['pwdHistoryLength'][0])

        password_properties = self.d2b(data[0]['attributes']['pwdProperties'][0])
        print("\t[+] Password Complexity Flags:", password_properties)
        print("")
        print("\t\t[+] Domain Refuse Password Change:", password_properties[0])
        print("\t\t[+] Domain Password Store Cleartext:", password_properties[1])
        print("\t\t[+] Domain Password Lockout Admins:", password_properties[2])
        print("\t\t[+] Domain Password No Clear Change:", password_properties[3])
        print("\t\t[+] Domain Password No Anon Change:", password_properties[4])
        print("\t\t[+] Domain Password Complex:", password_properties[5])
        print("")
        print("\t[+] Maximum password age:", self.convert(data[0]['attributes']['maxPwdAge'][0]))
        print("\t[+] Minimum password age:", self.convert(data[0]['attributes']['minPwdAge'][0]))
        print("\t[+] Reset Account Lockout Counter:", self.convert(data[0]['attributes']['lockoutDuration'][0]))
        print("\t[+] Account Lockout Threshold:", data[0]['attributes']['lockoutThreshold'][0])
        print("\t[+] Forced Log off Time:", self.convert(data[0]['attributes']['forceLogoff'][0]))

    def domain_info(self, data):
        print('''
    +--------------------------------------+
    | Getting Domain Sid For               |
    +--------------------------------------+
    ''')
        print('[+] Domain Name:', data[0]['attributes']['dc'][0])
        print('Domain Sid:', data[0]['attributes']['objectSid'][0])
        print('')
        return data[0]['attributes']['dc'][0]

    def user_info(self, users, dc):
        print('''
    +------------------------+
    | Users Infos            |
    +------------------------+
    ''')
        for user in users:
            desc = user['attributes'].get('description')[0] if user['attributes'].get('description') else "(null)"
            print("Account: " + dc + "\\" + user['attributes']['sAMAccountName'][0] + "\tName: " + user['attributes']['name'][0] + "\tDesc: " + desc)

        print("")
        for user in users:
            print("user:[" + user['attributes']['sAMAccountName'][0] + "]")
        print("")

    def groups_info(self, groups, dc):
        print('''
    +------------------------+import os.path
    | Groups Infos           |
    +------------------------+
    ''')
        for group in groups:
            print("group:[" + group['attributes']['name'][0] + "]")

        for group in groups:
            if group['attributes'].get('member'):
                users = re.findall(r"^CN=([\w\s\-\_\{\}\.\$\#]+)", '\n'.join(group['attributes']['member']), re.M)
                if users:
                    print("\n[+] Getting domain group memberships:")
                    for user in users:
                        if user == "S-1-5-11":
                            user = "NT AUTHORITY\Authenticated Users"
                        elif user == "S-1-5-4":
                            user = "NT AUTHORITY\INTERACTIVE"
                        elif user == "S-1-5-17":
                            user = "NT AUTHORITY\IUSR"
                        print("Group '" + group['attributes']['name'][0] + "' has member: " + dc + "\\" + user)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='LDAPDomainDump to pretty output like enum4linux.')

    #Main parameters
    parser.add_argument("-d", "--directory", help="The ldapdomaindump directory where the json files are saved. Required files: domain_users.json, domain_groups.json and domain_policy.json")
    args = parser.parse_args()

    if args.directory:
        with open(args.directory + '/domain_policy.json') as f:
            domain = json.load(f)
        with open(args.directory + '/domain_users.json') as f:
            users = json.load(f)
        with open(args.directory + '/domain_groups.json') as f:
            groups = json.load(f)
        pretty = PrettyOuput()
        dc = pretty.domain_info(domain)
        pretty.password_complexity(domain)
        pretty.user_info(users, dc.upper())
        pretty.groups_info(groups, dc.upper())
    else:
        print("Missing parameter --directory /output/")
<|endoftext|>