PY2 = sys.version_info[0] == 2
def with_metaclass(meta, *bases):
        def __new__(cls, name, this_bases, d):
def csv_open(filename, mode='r'):
    class metaclass(meta): def merge_arrays(array1, array2, result_fields='union', default_values=None):
    """
    data in array2 overrides data in array1
    both arrays must have 'id' fields
    """

    fields1 = get_fields(array1)
    fields2 = get_fields(array2)

    # TODO: check that common fields have the same type
    if result_fields == 'union':
        names1 = set(array1.dtype.names)
        fields_notin1 = [(name, type_) for name, type_ in fields2
                         if name not in names1]
        output_fields = fields1 + fields_notin1
    elif result_fields == 'array1':
        output_fields = fields1
    else:
        raise ValueError('%s in not a valid value for result_fields argument' %
                         result_fields)

    output_dtype = np.dtype(output_fields)

    ids1 = array1['id']
    ids2 = array2['id']
    all_ids = np.union1d(ids1, ids2)
    max_id = all_ids[-1]

    # compute new id_to_rownum
    id_to_rownum = np.full(max_id + 1, -1, dtype=int)
    for rownum, rowid in enumerate(all_ids):
        id_to_rownum[rowid] = rownum

    # 1) create resulting array
    ids1_complete = len(ids1) == len(all_ids)
    ids2_complete = len(ids2) == len(all_ids)
    output_is_arr1 = array1.dtype == output_dtype and ids1_complete
    output_is_arr2 = array2.dtype == output_dtype and ids2_complete
    arr1_complete = set(fields1) >= set(output_fields) and ids1_complete
    arr2_complete = set(fields2) >= set(output_fields) and ids2_complete
    if output_is_arr2:
        output_array = array2
    elif output_is_arr1:
        # TODO: modifying array1 in-place suits our particular needs for now
        # but it should really be a (non-default) option
        output_array = array1
    elif arr1_complete or arr2_complete:
        output_array = np.empty(len(all_ids), dtype=output_dtype)
    else:
        output_array = get_default_array(len(all_ids), output_dtype,
                                         default_values)

    # 2) copy data from array1 (if it will not be overridden)
    if not arr2_complete:
        output_array = merge_subset_in_array(output_array, id_to_rownum,
                                             array1, first=True,
                                             default_values=default_values)

    # 3) copy data from array2
    if not output_is_arr2:
        output_array = merge_subset_in_array(output_array, id_to_rownum, array2,
                                             default_values=default_values)

    return output_array, id_to_rownum def get_fields(array):
    dtype = array.dtype
    return [(name, normalize_type(dtype[name].type)) for name in dtype.names] class ColumnArray(object):
    def __init__(self, array=None):
        columns = {}
        if array is not None:
            if isinstance(array, (np.ndarray, ColumnArray)):
                for name in array.dtype.names:
                    columns[name] = array[name].copy()
                self.dtype = array.dtype
                self.columns = columns
            elif isinstance(array, list):
                for name, column in array:
                    columns[name] = column
                self.dtype = np.dtype([(name, column.dtype)
                                       for name, column in array])
                self.columns = columns
            else:
                # TODO: make a property instead?
                self.dtype = None
                self.columns = columns
        else:
            self.dtype = None
            self.columns = columns

    def __getitem__(self, key):
        if isinstance(key, basestring):
            return self.columns[key]
        else:
            # int, slice, ndarray
            ca = ColumnArray()
            for name, colvalue in self.columns.items():
                ca[name] = colvalue[key]
            ca.dtype = self.dtype
            return ca

    def __setitem__(self, key, value):
        """does not copy value except if a type conversion is necessary"""

        if isinstance(key, basestring):
            if isinstance(value, (np.ndarray, la.LArray)) and value.shape:
                column = value
            else:
                # expand scalars (like ndarray does) so that we don't have to
                # check isinstance(x, ndarray) and x.shape everywhere
                column = np.full(len(self), value, dtype=gettype(value))

            if key in self.columns:
                # converting to existing dtype
                if column.dtype != self.dtype[key]:
                    column = column.astype(self.dtype[key])
                self.columns[key] = column
            else:
                # adding a new column so we need to update the dtype
                self.columns[key] = column
                self._update_dtype()

#            ids = {}
#            for k, v in self.columns.items():
#                ids.setdefault(id(v), set()).add(k)
#            dupes = [v for k, v in ids.items() if len(v) > 1]
#            if dupes:
#                print "aliases", dupes
        else:
            # int, slice, ndarray
            for name, column in self.columns.items():
                column[key] = value[name]

    def put(self, indices, values, mode='raise'):
        for name, column in self.columns.items():
            column.put(indices, values[name], mode)

    @property
    def nbytes(self):
        return sum(v.nbytes for v in self.columns.values())

    def __delitem__(self, key):
        del self.columns[key]
        self._update_dtype()

    def _update_dtype(self):
        # handle fields already present (iterate over old dtype to preserve
        # order)
        if self.dtype is not None:
            old_fields = self.dtype.names
            fields = [(name, self.dtype[name])
                      for name in old_fields
                      if name in self.columns]
        else:
            old_fields = []
            fields = []
        # add new fields (not already handled)
        old_fields = set(old_fields)
        fields += [(name, column.dtype)
                   for name, column in self.columns.items()
                   if name not in old_fields]
        self.dtype = np.dtype(fields)

    def __len__(self):
        if len(self.columns):
            anycol = next(iter(self.columns.values()))
            return len(anycol)
        else:
            return 0

    def keep(self, key):
        """key can be either a vector of int indices or boolean filter"""

        # using gc.collect() after each column update frees a bit of memory
        # but slows things down significantly.
        for name, column in self.columns.items():
            self.columns[name] = column[key]

    def append(self, array):
        assert array.dtype == self.dtype, (array.dtype, self.dtype)
        # using gc.collect() after each column update frees a bit of memory
        # but slows things down significantly.
        for name, column in self.columns.items():
            self.columns[name] = np.concatenate((column, array[name]))

    def append_to_table(self, table, buffersize=10 * MB):
        append_carray_to_table(self, table, buffersize=buffersize)

    @classmethod
    def empty(cls, length, dtype):
        ca = cls()
        for name in dtype.names:
            ca.columns[name] = np.empty(length, dtype[name])
        ca.dtype = dtype
        return ca

    @classmethod
    def from_table(cls, table, start=0, stop=None, buffersize=10 * MB):
        # reading a table one column at a time is very slow, this is why this
        # function is even necessary
        if stop is None:
            stop = len(table)
        dtype = table.dtype
        max_buffer_rows = buffersize // dtype.itemsize
        numlines = stop - start
        ca = cls.empty(numlines, dtype)
        buffer_rows = min(numlines, max_buffer_rows)
        chunk = np.empty(buffer_rows, dtype=dtype)
        array_start = 0
        table_start = start
        while numlines > 0:
            buffer_rows = min(numlines, max_buffer_rows)
            if buffer_rows < len(chunk):
                # last chunk is smaller
                chunk = np.empty(buffer_rows, dtype=dtype)
            table.read(table_start, table_start + buffer_rows, out=chunk)
            ca[array_start:array_start + buffer_rows] = chunk
            table_start += buffer_rows
            array_start += buffer_rows
            numlines -= buffer_rows
        return ca

    @classmethod
    def from_table_coords(cls, table, indices, buffersize=10 * MB):
        dtype = table.dtype
        max_buffer_rows = buffersize // dtype.itemsize
        numlines = len(indices)
        ca = cls.empty(numlines, dtype)
        buffer_rows = min(numlines, max_buffer_rows)
        # chunk = np.empty(buffer_rows, dtype=dtype)
        start, stop = 0, buffer_rows
        while numlines > 0:
            buffer_rows = min(numlines, max_buffer_rows)
            # if buffer_rows < len(chunk):
            #    # last chunk is smaller
            #    chunk = np.empty(buffer_rows, dtype=dtype)
            chunk_indices = indices[start:stop]
            # as of PyTables 3.2.2, read_coordinates does not support out=
            # table.read_coordinates(chunk_indices, out=chunk)
            chunk = table.read_coordinates(chunk_indices)
            ca[start:stop] = chunk
            start += buffer_rows
            stop += buffer_rows
            numlines -= buffer_rows
        return ca

    def add_and_drop_fields(self, names_to_keep, output_fields, default_values):
        """modify inplace.
        Only passing output_fields is not enough because one may want to reset a field data (see issue 227).
        """

        output_dtype = np.dtype(output_fields)
        output_names = set(output_dtype.names)
        input_names = set(self.dtype.names)
        # drop extra fields
        for name in input_names - set(names_to_keep):
            del self[name]

        # add missing fields
        length = len(self)
        if default_values is None:
            default_values = {}
        for name in output_names - set(self.dtype.names):
            self[name] = get_default_vector(length, output_dtype[name], default_values[name]) def index_table(table):
    """
    table is an iterable of rows, each row is a mapping (name -> value).
    Rows must contain at least 'period' and 'id' columns and must be sorted
    by period.
    """
    rows_per_period = {}
    id_to_rownum_per_period = {}
    temp_id_to_rownum = []
    max_id_so_far = -1
    current_period = None
    start_row = None
    for idx, row in enumerate(table):
        period, row_id = row['period'], row['id']
        if period != current_period:
            if current_period is not None and period < current_period:
                msg = "data is not ordered by period " \
                      "({} at data line {} is < {})"
                raise Exception(msg.format(period, idx + 1, current_period))
            if start_row is not None:
                rows_per_period[current_period] = start_row, idx
                # assumes the data is sorted on period then id
                id_to_rownum = np.array(temp_id_to_rownum)
                id_to_rownum_per_period[current_period] = id_to_rownum
                temp_id_to_rownum = [-1] * (max_id_so_far + 1)
            start_row = idx
            current_period = period
        if row_id > max_id_so_far:
            extra = [-1] * (row_id - max_id_so_far)
            temp_id_to_rownum.extend(extra)
        if temp_id_to_rownum[row_id] != -1:
            msg = "duplicate row for id {} for period {} (at data line {})"
            # idx + 1 is correct for ViTables, which starts counting at 1, but
            # is still off by one (or more) for .csv files because of headers
            # and comments
            raise Exception(msg.format(row_id, period, idx + 1))
        temp_id_to_rownum[row_id] = idx - start_row
        max_id_so_far = max(max_id_so_far, row_id)
    if current_period is not None:
        rows_per_period[current_period] = (start_row, len(table))
        id_to_rownum_per_period[current_period] = np.array(temp_id_to_rownum)
    return rows_per_period, id_to_rownum_per_period def build_period_array(input_table, fields_to_keep, output_fields, input_rows,
                       input_index, start_period, default_values):
    periods_before = [p for p in input_rows.keys() if p <= start_period]
    if not periods_before:
        id_to_rownum = np.empty(0, dtype=int)
        output_array = ColumnArray.empty(0, np.dtype(output_fields))
        return output_array, id_to_rownum

    periods_before.sort()
    # take the last period which we have data for
    target_period = periods_before[-1]

    # computing is_present
    max_id = len(input_index[target_period]) - 1
    period_id_to_rownum = None
    present_in_period = None
    is_present = np.zeros(max_id + 1, dtype=bool)
    for period in periods_before:
        period_id_to_rownum = input_index[period]
        present_in_period = period_id_to_rownum != -1
        present_in_period.resize(max_id + 1, refcheck=False)
        is_present |= present_in_period

    # if all individuals are present in the target period, we are done already!
    if np.array_equal(present_in_period, is_present):
        start, stop = input_rows[target_period]
        input_array = ColumnArray.from_table(input_table, start, stop)
        input_array.add_and_drop_fields(fields_to_keep, output_fields, default_values)
        return input_array, period_id_to_rownum

    # building id_to_rownum for the target period
    id_to_rownum = np.full(max_id + 1, -1, dtype=int)
    rownum = 0
    for row_id, present in enumerate(is_present):
        if present:
            id_to_rownum[row_id] = rownum
            rownum += 1

    # computing the source row for each destination row
    # we loop over the periods before start_period in reverse order
    output_array_source_rows = np.full(rownum, -1, dtype=int)
    for period in periods_before[::-1]:
        start, stop = input_rows[period]
        input_rownums = np.arange(start, stop)

        input_id_to_rownum = input_index[period]
        id_is_in_period = input_id_to_rownum != -1

        # which output rows are filled by input for this period
        output_rownums = id_to_rownum[np.where(id_is_in_period)]

        # get source rows (in the global array) for individuals in this period
        source_rows = output_array_source_rows[output_rownums]

        # if their source row is already known, leave them alone
        need_update = source_rows == -1

        # global indices of rows which are not set yet (for this period)
        rows_to_update = output_rownums[need_update]

        # source row for those rows
        local_source_rows = input_rownums[need_update]

        # update the source row for those rows
        safe_put(output_array_source_rows, rows_to_update, local_source_rows)

        if np.all(output_array_source_rows != -1):
            break

    # reading data
    output_array = ColumnArray.from_table_coords(input_table,
                                                 output_array_source_rows)
    output_array.add_and_drop_fields(fields_to_keep, output_fields, default_values)
    return output_array, id_to_rownum class Variable(Expr):
    __children__ = ()

    def __init__(self, entity, name, dtype=None):
        # from entities import Entity
        # assert entity is None or isinstance(entity, Entity)
        self.entity = entity
        self.name = name
        self._dtype = dtype
        # self.version = 0
        # self.used = 0

    def __repr__(self):
        return "%s.%s" % (self.entity, self.name)

    def __str__(self):
        return self.name
    as_string = __str__

    def as_simple_expr(self, context):
        return self

    def dtype(self, context):
        if self._dtype is None and self.name in context:
            return gettype(context[self.name])
        else:
            return self._dtype class VariableMethodHybrid(Variable):
    def __call__(self, *args, **kwargs):
        return MethodCall(self.entity, self.name, args, kwargs) class GlobalVariable(EvaluableExpression):
    __children__ = ()

    def __init__(self, tablename, name, dtype):
        self.tablename = tablename
        self.name = name
        self._dtype = dtype

    def __repr__(self):
        if self.name is None:
            return self.tablename
        else:
            return "%s.%s" % (self.tablename, self.name)

    def get_tmp_varname(self, context):
        period = self._eval_key(context)
        if isinstance(period, int):
            return '__%s_%s_%s' % (self.tablename, self.name, period)
        else:
            return EvaluableExpression.get_tmp_varname(self, context)

    def _eval_key(self, context):
        return context.period

    def evaluate(self, context):
        globals_data = context.global_tables
        globals_table = globals_data[self.tablename]
        if self.name is None:
            return globals_table

        if self.name not in globals_table.dtype.fields:
            raise Exception("Unknown global: %s" % self.name)

        key = self._eval_key(context)
        # TODO: this row computation should be encapsulated in the
        # globals_table object and the index column should be configurable
        colnames = globals_table.dtype.names
        if 'period' in colnames or 'PERIOD' in colnames:
            try:
                globals_periods = globals_table['PERIOD']
            except ValueError:
                globals_periods = globals_table['period']
            base_period = globals_periods[0]
            if isinstance(key, slice):
                translated_key = slice(key.start - base_period,
                                       key.stop - base_period,
                                       key.step)
            else:
                translated_key = key - base_period
        else:
            translated_key = key

        column = globals_table[self.name]
        numrows = len(column)
        missing_value = get_default_value(column)

        if isinstance(translated_key, np.ndarray) and translated_key.shape:
            return safe_take(column, translated_key, missing_value)
        elif isinstance(translated_key, slice):
            start, stop = translated_key.start, translated_key.stop
            step = translated_key.step
            if step is not None and step != 1:
                raise NotImplementedError("step != 1 (%d)" % step)
            if (isinstance(start, np.ndarray) and start.shape or
                    isinstance(stop, np.ndarray) and stop.shape):
                lengths = stop - start
                length0 = lengths[0]
                if not isinstance(start, np.ndarray) or not start.shape:
                    start = np.repeat(start, len(lengths))
                if np.all(lengths == length0):
                    # constant length => result is a 2D array:
                    # num_individuals x slice_length
                    result = np.empty((len(lengths), length0),
                                      dtype=column.dtype)
                    # we assume there are more individuals than there are
                    # "periods" (or other ticks) in the table.
                    # XXX: We might want to actually test that it is true and
                    # loop on the individuals instead if that is not the case
                    for i in range(length0):
                        result[:, i] = safe_take(column, start + i,
                                                 missing_value)
                    return result
                else:
                    # varying length => result is an array (num_individuals) of
                    # 1D arrays (slice lengths)
                    # each "item" of the result is a view, so we pay "only" for
                    # all the arrays overhead, not for the data itself.
                    result = np.empty(len(lengths), dtype=object)
                    if not isinstance(stop, np.ndarray) or not stop.shape:
                        stop = np.repeat(stop, len(lengths))
                    for i in range(len(lengths)):
                        result[i] = column[start[i]:stop[i]]
                    return IrregularNDArray(result)
            else:
                # out of bounds slices bounds are "dropped" silently (like in
                # python) -- ie the length of the slice returned can be
                # smaller than the one asked. We could return "missing_value"
                # for indices out of bounds but I do not know if it would be
                # better. Since this version is easier to implement, lets go for
                # it for now.
                return column[translated_key]
        else:
            # scalar key
            out_of_bounds = (translated_key < 0) or (translated_key >= numrows)
            return column[translated_key] if not out_of_bounds \
                else missing_value

    def __getitem__(self, key):
        return SubscriptedGlobal(self.tablename, self.name, key, self._dtype)

    def dtype(self, context):
        return self._dtype class GlobalTable(object):
    def __init__(self, name, fields):
        """fields is a list of tuples (name, type)"""

        self.name = name
        self.fields = fields
        self.fields_map = dict(fields)

    def __getattr__(self, key):
        return GlobalVariable(self.name, key, self.fields_map[key])

    # noinspection PyUnusedLocal
    def traverse(self, context):
        yield self

    def __repr__(self):
        # Remember this is the expression (only used via qshow, ...), so we do
        # not want to print the data in here
        return 'Table(%s)' % ', '.join([name for name, _ in self.fields]) class GlobalArray(Variable):
    def __init__(self, name, dtype=None, autoindex=None):
        Variable.__init__(self, None, name, dtype)
        # convert to tuple so that it is hashable
        if isinstance(autoindex, list):
            autoindex = tuple(autoindex)
        self.autoindex = autoindex

    def as_simple_expr(self, context):
        globals_data = context.global_tables
        result = globals_data[self.name]
        if self.autoindex is not None:
            result = index_array_by_variables(result, context, result.axes[self.autoindex])
        # XXX: maybe I should just use self.name?
        # FIXME: use self.add_tmp_var, because in combination with autoindex,
        # the variable could have a different value
        tmp_varname = '__%s' % self.name
        if tmp_varname in context:
            array = context[tmp_varname]
            assert isinstance(array, la.LArray)
            assert context[tmp_varname].equals(result)
        context[tmp_varname] = result
        return Variable(context.entity, tmp_varname) class Expr(object):
    # XXX: I wonder if those couldn't be computed automatically by using
    # isinstance(v, Expr)
    __children__ = ()
    num_tmp = 0

    def __init__(self):
        raise NotImplementedError()

    @property
    def children(self):
        d = self.__dict__
        return tuple(d[k] for k in self.__children__)

    @property
    def value(self):
        d = self.__dict__
        # ignore _ keys, for example "_variables"
        keys = sorted(k for k in d.keys() if not k.startswith('_'))
        children = set(self.__children__)
        return tuple(d[k] for k in keys if k not in children)

    # makes sure we do not use "normal" python logical operators
    # (and, or, not)
    def __nonzero__(self):
        raise Exception("Improper use of boolean operators, you probably "
                        "forgot parenthesis around operands of an 'and' or "
                        "'or' expression. The complete expression cannot be "
                        "displayed but it contains: '%s'." % str(self))

    def evaluate(self, context):
        # period = context.period
        #
        # if isinstance(period, np.ndarray):
        #     assert np.isscalar(period) or not period.shape
        #     period = int(period)

        # cache_key = (self, period, context.entity_name, context.filter_expr)
        # try:
        #     cached_result = expr_cache.get(cache_key, None)
        #     # FIXME: lifecycle functions should invalidate all variables!
        #     if cached_result is not None:
        #         return cached_result
        # except TypeError:
        #     # The cache_key failed to hash properly, so the expr is not
        #     # cacheable. It *should* be because of a not_hashable expr
        #     # somewhere within cache_key[3].
        #     cache_key = None

        simple_expr = self.as_simple_expr(context)
        if isinstance(simple_expr, Variable) and simple_expr.name in context:
            return context[simple_expr.name]

        # check for labeled arrays, to work around the fact that numexpr
        # does not preserve ndarray subclasses.

        # avoid checking for arrays types in the past, because that is a
        # costly operation (context[var_name] fetches the column from disk
        # in that case). This probably prevents us from doing stuff like
        # align(lag(groupby() / groupby())), but it is a limitation I can
        # live with to avoid hitting the disk twice for each disk access.

        # TODO: I should rewrite this whole mess when my "dtype" method
        # supports ndarrays and la.LArray so that I can get the dtype from
        # the expression instead of from actual values.
        expr_axes = None
        numexpr_eval = True
        assert isinstance(context, EvaluationContext), type(context)
        local_ctx = context.entity_data
        if isinstance(local_ctx, EntityContext) and local_ctx.is_array_period:
            for var in simple_expr.collect_variables():
                assert var.entity is None or var.entity is context.entity, \
                    "should not have happened (as_simple_expr should " \
                    "have transformed non-local variables)"

                # var_name should always be in the context at this point
                # because missing temporaries should have been already caught
                # in expr_eval
                value = context[var.name]
                # value = local_ctx[var.name]

                # check that LArrays (if any) have all the same axes and bypass numexpr otherwise
                if isinstance(value, la.LArray):
                    if expr_axes is None:
                        expr_axes = value.axes
                    else:
                        if value.axes != expr_axes:
                            numexpr_eval = False

        # TODO: when numexpr_eval is False, we should bypass the string roundtrip
        s = simple_expr.as_string()
        constants = {'nan': float('nan'), 'inf': float('inf')}
        if numexpr_eval:
            try:
                res = evaluate(s, local_ctx, constants, truediv='auto')
                if expr_axes is not None:
                    # This relies on the fact that currently all the expression we evaluate through numexpr preserve
                    # array shapes, but if we ever use numexpr reduction capabilities, we will be in trouble
                    res = la.LArray(res, expr_axes)
            except Exception:
                if debug:
                    print("evaluate failed")
                    print("s:", s)
                    print("local context:", local_ctx)
                    print("constants:", constants)
                raise
        else:
            res = eval(s, constants, local_ctx)

        if isinstance(res, np.ndarray) and not res.shape:
            res = np.asscalar(res)

        # if cache_key is not None:
        #     expr_cache[cache_key] = res
        #     if cached_result is not None:
        #         assert np.array_equal(res, cached_result), \
        #             "%s != %s" % (res, cached_result)
        return res

    def as_simple_expr(self, context):
        """
        evaluate any construct that is not supported by numexpr and
        create temporary variables for them
        """
        raise NotImplementedError()

    def as_string(self):
        raise NotImplementedError()

    def __getitem__(self, key):
        # TODO: we should be able to know at "compile" time if this is a
        # scalar or a vector and disallow getitem in case of a scalar
        return SubscriptedExpr(self, key)

    def __getattr__(self, key):
        if key == '_variables':
            raise AttributeError("%s (of type '%s') has no attribute '%s'"
                                     % (self, self.__class__.__name__, key))
        else:
            return ExprAttribute(self, key)
        # if key in {'data', 'dtype', 'itemsize', 'nbytes', 'ndim', 'shape', 'size',
        #            'axes', 'points', 'row_totals', 'col_totals',
        #            # aggregates
        #            'all', 'any', 'max', 'mean', 'min', 'prod', 'ptp', 'std',
        #            'sum', 'var', 'cumprod', 'cumsum',
        #            # element-wise
        #            'astype', 'clip', 'copy', 'round',
        #            # inplace
        #            'fill', 'partition', 'sort',
        #            # indirect
        #            'argmax', 'argmin', 'argpartition', 'argsort',
        #            # other
        #            'nonzero', 'reshape', 'transpose',
        #            # should go away when I implement len()
        #            '__len__'}:
        #     # excluded (and it would take some convincing to add them):
        #     # base, byteswap, ctypes, data, dump, dumps, getfield, item,
        #     # itemset, newbyteorder, put, resize, setfield, setflags, swapaxes,
        #     # take, trace, view
        #
        #     # excluded for now (I am open if anybody asks for them):
        #     # choose, compress, conj, conjugate, diagonal, dot, flags, flat,
        #     # flatten, imag, ravel, real, repeat, squeeze, strides
        #
        #     # compress (select using a boolean index for one axis) is nice but
        #     # I would like to implement a[bool_idx_wh_axis] instead in LArray
        #     return ExprAttribute(self, key)
        # else:
        #     raise AttributeError("%s (of type '%s') has no attribute '%s'"
        #                          % (self, self.__class__.__name__, key))

    def traverse(self):
        for child in self.children:
            for node in traverse_expr(child):
                yield node
        yield self

    def all_of(self, node_type):
        """
        Yields nodes corresponding to node_type.

        Parameters
        ----------
        node_type : class or tuple of class

        Yields
        ------
        all nodes of the tree (possibly including self) corresponding to
        node_type
        """
        for node in self.traverse():
            if isinstance(node, node_type):
                yield node

    def collect_variables(self):
        # TODO: it would be cleaner if we initialized _variables in __init__,
        # however it means each Expr subclass would have to call its parent
        # __init__ (which is a good thing but too much hassle at this point).
        if not hasattr(self, "_variables") or self._variables is None:
            # FIXME: this is a quick hack to make "othertable" work.
            # We should return prefixed variable instead.
            badvar = lambda v: isinstance(v, ShortLivedVariable) or \
                               (isinstance(v, GlobalVariable) and
                                v.tablename != 'periodic')
            self._variables = set(v for v in self.all_of(Variable)
                                  if not badvar(v))
        return self._variables

    # TODO: make equivalent/commutative expressions compare equal and hash to
    # the same thing.
    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return False

        def strict_equal(a, b):
            # we are intentionally very picky so that eg. True != 1
            if type(a) is not type(b):
                return False
            if isinstance(a, (tuple, list)):
                return all(strict_equal(e1, e2) for e1, e2 in zip(a, b))
            else:
                return a == b

        res = strict_equal(self.value, other.value) and \
            strict_equal(self.children, other.children)
        if res:
            if str(self) != str(other):
                print()
                print('SHOULD NOT COMPARE EQUAL!')
                print(str(self).ljust(40), '>>>', self.value, self.children)
                print(str(other).ljust(40), '>>>', other.value, other.children)
                print(type(self), type(other),
                      len(self.children), len(other.children))
                print([(x, type(x)) for x in self.children])
                print([(x, type(x)) for x in other.children])
                # raise Exception("should not compare equal")
        return res

    def __hash__(self):
        return hash((self.__class__.__name__, self.value,
                     make_hashable(self.children)))

    def __contains__(self, expr):
        for node in self.traverse():
            if expr == node:
                return True
        return False

    def get_tmp_varname(self, context):
        tmp_varname = "temp_%d" % self.num_tmp
        Expr.num_tmp += 1
        return tmp_varname

    def add_tmp_var(self, context, result):
        tmp_varname = self.get_tmp_varname(context)
        if tmp_varname in context:
            # should be consistent but nan != nan
            if isinstance(result, np.ndarray):
                assert array_nan_equal(context[tmp_varname], result)
            else:
                assert result != result or context[tmp_varname] == result
        # FIXME: we should never modify the context in-place. We should rather
        #        have a build_context method.
        context[tmp_varname] = result
        return Variable(context.entity, tmp_varname, gettype(result)) class BinaryOp(Expr):
    __children__ = ('expr1', 'expr2')

    def __init__(self, op, expr1, expr2):
        self.op = op
        self.expr1 = expr1
        self.expr2 = expr2

    def as_simple_expr(self, context):
        expr1 = as_simple_expr(self.expr1, context)
        expr2 = as_simple_expr(self.expr2, context)
        return self.__class__(self.op, expr1, expr2)

    # We can't simply use __str__ because of where vs if
    def as_string(self):
        expr1, expr2 = as_string(self.expr1), as_string(self.expr2)
        return "(%s %s %s)" % (expr1, self.op, expr2)

    def dtype(self, context):
        return coerce_types(context, self.expr1, self.expr2)

    # FIXME: only add parentheses if necessary
    def __repr__(self):
        nicerop = {'&': 'and', '|': 'or'}
        niceop = nicerop.get(self.op, self.op)
        return "(%r %s %r)" % (self.expr1, niceop, self.expr2) class MethodSymbol(object):
    def __init__(self, name, entity):
        self.name = name
        self.entity = entity

    def __call__(self, *args, **kwargs):
        # we cannot use self.entity.processes as they are not defined yet (we
        # are probably currently building them), so we cannot return a
        # GenericFunctionCall now like we should and instead must either
        # return an intermediary object (MethodCallToResolve) which we will
        # "resolve" later, or use DynamicFunctionCall (for which we cannot
        # have dtype yet). However that "resolve" step is currently hard to
        # do because we need ast.NodeTransformer-like machinery which we do not
        # have yet.
        # return MethodCallToResolve(self.entity, self.name, args, kwargs)
        return MethodCall(self.entity, self.name, args, kwargs)

    def __repr__(self):
        return 'MethodSymbol({}, {})'.format(self.name, self.entity)

    def __str__(self):
        return '{}.{}'.format(self.entity, self.name) def normalize_type(type_):
    return idx_to_type[type_to_idx[type_]] def parse(s, context, interactive=False):
    globals_context = {'False': False,
                       'True': True,
                       'nan': float('nan'),
                       'inf': float('inf'),
                       'float': float,
                       'int': int,
                       'bool': bool}
    globals_context.update(functions)
    globals_context.update(context.get('__globals__', {}))
    # modify in-place
    context['__globals__'] = globals_context
    try:
        node = _parse(s, interactive=interactive)
        return to_ast(node, context)
    except Exception as e:
        add_context(e, "while parsing: " + s)
        raise class Assignment(Process):
    def __init__(self, name, entity, expr):
        super(Assignment, self).__init__(name, entity)
        self.expr = expr
        self.temporary = name not in entity.fields.names

    def run(self, context):
        value = expr_eval(self.expr, context)
        # Assignment to a field with a name == None is valid: it simply means
        # the result must not be stored. This happens when a user does not
        # store anywhere the result of an expression (it usually has side
        # effects -- csv, new, remove, ...).
        if self.name is not None:
            self.store_result(value, context)

    def store_result(self, result, context):
        if isinstance(result, (np.ndarray, la.LArray)):
            res_type = result.dtype.type
        else:
            res_type = type(result)

        if self.temporary:
            target = self.entity.temp_variables
        else:
            # we cannot store/cache self.entity.array[self.name] because the
            # array object can change (eg when enlarging it due to births)
            target = self.entity.array
            result = np.asarray(result)
            # TODO: assert type for temporary variables too
            target_type_idx = type_to_idx[target[self.name].dtype.type]
            res_type_idx = type_to_idx[res_type]
            if res_type_idx > target_type_idx:
                raise Exception(
                    "trying to store %s value into '%s' field which is of "
                    "type %s" % (idx_to_type[res_type_idx].__name__,
                                 self.name,
                                 idx_to_type[target_type_idx].__name__))

        # the whole column is updated
        target[self.name] = result

        # invalidate cache
        period = context.period
        if isinstance(period, np.ndarray):
            assert np.isscalar(period) or not period.shape
            period = int(period)
        expr_cache.invalidate(period, context.entity_name,
                              Variable(self.entity, self.name))

    def expressions(self):
        if isinstance(self.expr, Expr):
            yield self.expr
 class ProcessGroup(Process):
    def __init__(self, name, entity, subprocesses, purge=True):
        super(ProcessGroup, self).__init__(name, entity)
        self.subprocesses = subprocesses
        self.calls = collections.Counter()
        self.purge = purge
        self.versions = {}

    def run_guarded(self, context):
        period = context.period

        if config.log_level == "processes":
            print()

        try:
            for k, v in self.subprocesses:
                if config.log_level == "processes":
                    print("    *", end=' ')
                    if k is not None:
                        print(k, end=' ')
                    utils.timed(v.run_guarded, context)
                else:
                    v.run_guarded(context)
                context.simulation.start_console(context)
        finally:
            if config.autodump is not None:
                self._autodump(context)

            if config.autodiff is not None:
                self._autodiff(period)

            if self.purge:
                self.entity.purge_locals()

    @property
    def predictors(self):
        return [v.name for _, v in self.subprocesses
                if isinstance(v, Assignment) and v.name is not None]

    @property
    def _modified_fields(self):
        fnames = self.predictors
        if not fnames:
            return []

        fnames.insert(0, 'id')
        temp = self.entity.temp_variables
        array = self.entity.array
        length = len(array)

        fields = [(k, temp[k] if k in temp else array[k])
                  for k in utils.unique(fnames)]
        return [(k, v) for k, v in fields
                if isinstance(v, np.ndarray) and v.shape == (length,)]

    def _tablename(self, period):
        self.calls[(period, self.name)] += 1
        num_calls = self.calls[(period, self.name)]
        if num_calls > 1:
            return '{}_{}'.format(self.name, num_calls)
        else:
            return self.name

    def _autodump(self, context):
        fields = self._modified_fields
        if not fields:
            return

        period = context.period
        fname, numrows = config.autodump
        h5file = config.autodump_file
        name = self._tablename(period)
        dtype = np.dtype([(k, v.dtype) for k, v in fields])
        table = h5file.create_table('/{}'.format(period), name, dtype,
                                    createparents=True)

        fnames = [k for k, _ in fields]
        print("writing {} to {}/{}/{} ...".format(', '.join(fnames),
                                                  fname, period, name))

        entity_context = EntityContext(context, self.entity, {'period': period})
        append_carray_to_table(entity_context, table, numrows)
        print("done.")

    def _autodiff(self, period, showdiffs=10, raiseondiff=False):
        fields = self._modified_fields
        if not fields:
            return

        fname, numrows = config.autodiff
        h5file = config.autodump_file
        tablepath = '/p{}/{}'.format(period, self._tablename(period))
        print("comparing with {}{} ...".format(fname, tablepath))
        if tablepath in h5file:
            table = h5file.getNode(tablepath)
            disk_array = ColumnArray.from_table(table, stop=numrows)
            diff_array(disk_array, ColumnArray(fields), showdiffs, raiseondiff)
        else:
            print("  SKIPPED (could not find table)")

    def expressions(self):
        for _, p in self.subprocesses:
            for e in p.expressions():
                yield e

    def ssa(self, fields_versions):
        function_vars = set(k for k, p in self.subprocesses if k is not None)
        global_vars = set(self.entity.variables.keys())
        local_vars = function_vars - global_vars

        local_versions = collections.defaultdict(int)
        for k, p in self.subprocesses:
            # mark all variables in the expression with their current version
            for expr in p.expressions():
                for node in expr.all_of(Variable):
                    versions = (local_versions if node.name in local_vars
                                else fields_versions)
                    # FIXME: for .version to be meaningful, I need to have
                    # a different variable instance each time the variable
                    # is used.
                    # >>> the best solution AFAIK is to parse the expressions
                    # in the same order as the "agespine".
                    # That way we will be able to type all temporary variables
                    # directly, and it would also solve the conditional
                    # context hack. There is no problem with temporary variables
                    # having different types over their lifetimes as these
                    # will actually be different variables (because their
                    # version will be different). There is no problem with
                    # a variable being different in two control flow
                    # "branches", because we do not have that case: if() coerce
                    # types.
                    # note that even if a branch is never "taken" (an if
                    # condition that is always True or always False or a forloop
                    # without any iteration), the type of the expressions
                    # that variables are assigned to in that branch will
                    # influence the type of the variable in subsequent code.
                    # XXX: what if I have a user-defined function that
                    # I call from two different places with an argument of a
                    # different type? ideally, it should generate two distinct
                    # functions, but I am not there yet. Having a check on the
                    # second call that the argument passed is of the same type
                    # than the signature type (which was inferred from the
                    # first call) seems enough for now.
                    node.version = versions[node.name]
                    node.used += 1
            # on assignment, increase the variable version
            if isinstance(p, Assignment):
                target = p.name
                versions = (local_versions if target in local_vars
                            else fields_versions)
                versions[target] += 1
 class While(Process):
    """this class implements while loops"""

    def __init__(self, name, entity, cond, code):
        """
        cond -- an Expr returning a (single) boolean, it means the condition
                value must be the same for all individuals
        code -- a ProcessGroup
        """
        Process.__init__(self, name, entity)
        self.cond = cond
        assert isinstance(code, ProcessGroup)
        self.code = code

    def run_guarded(self, context):
        while expr_eval(self.cond, context):
            self.code.run_guarded(context)
            # FIXME: this is a bit brutal :) This is necessary because
            # otherwise test_while loops indefinitely (because "values" is
            # never incremented)
            expr_cache.clear()

    def expressions(self):
        if isinstance(self.cond, Expr):
            yield self.cond
        for e in self.code.expressions():
            yield e
 class Function(Process):
    """this class implements user-defined functions"""

    def __init__(self, name, entity, argnames, code=None):
        """
        args -- a list of strings
        code -- a ProcessGroup (or None)
        """
        Process.__init__(self, name, entity)

        assert isinstance(argnames, list)
        assert all(isinstance(a, basestring) for a in argnames)
        self.argnames = argnames

        assert code is None or isinstance(code, ProcessGroup)
        self.code = code

    def run_guarded(self, context, *args, **kwargs):
        # XXX: wouldn't some form of cascading context make all this junk much
        # cleaner? Context(globalvars, localvars) (globalvars contain both
        # entity fields and global temporaries)

        backup = self.backup_and_purge_locals()
        # in case the function call is in the middle of a larger expression (see issue #186)
        backup_extra = context.entity_data.extra

        if len(args) != len(self.argnames):
            raise TypeError("%s() takes exactly %d arguments (%d given)" %
                            (self.name, len(self.argnames), len(args)))

        # TODO: this should be done in a separate "validate" phase. Unsure we can reliably do it in __init__ in
        # case we decide to parse methods before fields for some reason.
        for name in self.argnames:
            if name in self.entity.fields.names:
                raise ValueError("function '%s' cannot have an argument named "
                                 "'%s' because there is a field with the "
                                 "same name" % (self.name, name))

        # contextual filter should not transfer to the called function (even
        # if that would somewhat make sense) because in many cases the
        # columns used in the contextual filter are not available within the
        # called function. This is only relevant for functions called within
        # an if() expression.
        context = context.clone(filter_expr=None)

        # add arguments to the local namespace
        for name, value in zip(self.argnames, args):
            # backup the variable if it existed in the caller namespace
            if name in self.entity.temp_variables:
                # we can safely assign to backup without checking if that name
                # was already assigned because it is not possible for a variable
                # to be both in entity.temp_variables and in backup (they are
                # removed from entity.temp_variables).
                backup[name] = self.entity.temp_variables.pop(name)

            # cannot use context[name] = value because that would store the
            # value in .extra, which is wiped at the start of each process
            # and we need it to be available across all processes of the
            # function
            self.entity.temp_variables[name] = value
        try:
            self.code.run_guarded(context)
            result = None
        except ReturnException as r:
            result = r.result
        self.purge_and_restore_locals(backup)
        # restore extra
        context.entity_data.extra = backup_extra
        return result

    def expressions(self):
        if self.code is not None:
            for e in self.code.expressions():
                yield e

    def backup_and_purge_locals(self):
        # backup and purge local variables
        backup = {}
        for name in self.entity.local_var_names:
            backup[name] = self.entity.temp_variables.pop(name)
        return backup

    def purge_and_restore_locals(self, backup):
        # purge the local from the function we just ran
        self.entity.purge_locals()
        # restore local variables for our caller
        for k, v in backup.items():
            self.entity.temp_variables[k] = v
 class Return(Process):
    def __init__(self, name, entity, result_expr):
        super(Return, self).__init__(name, entity)
        self.result_expr = result_expr

    def run_guarded(self, context):
        raise ReturnException(expr_eval(self.result_expr, context))

    def expressions(self):
        if isinstance(self.result_expr, Expr):
            yield self.result_expr
 def count_occurrences(seq):
    counter = defaultdict(int)
    for e in seq:
        counter[e] += 1
    return list(counter.items()) def field_str_to_type(str_type, context):
    """
    Converts a (field) string type to its Python type.
    """
    if str_type not in str_to_type:
        raise SyntaxError("'%s' is not a valid type for %s."
                          % (str_type, context))
    return str_to_type[str_type] def size2str(value):
    """
    >>> size2str(0)
    '0 bytes'
    >>> size2str(100)
    '100 bytes'
    >>> size2str(1023)
    '1023 bytes'
    >>> size2str(1024)
    '1.00 Kb'
    >>> size2str(2000)
    '1.95 Kb'
    >>> size2str(10000000)
    '9.54 Mb'
    >>> size2str(1.27 * 1024 ** 3)
    '1.27 Gb'
    """
    units = ["bytes", "Kb", "Mb", "Gb", "Tb", "Pb"]
    scale = int(math.log(value, 1024)) if value else 0
    fmt = "%.2f %s" if scale else "%d %s"
    return fmt % (value / 1024.0 ** scale, units[scale]) class WarnOverrideDict(dict):
    def update(self, other=None, **kwargs):
        # copy the items to not lose them in case it is an exhaustible
        # iterable
        # also converts list and tuple to dict
        if not isinstance(other, dict):
            other = dict(other)
        self._intersect_warn(self, other)
        self._intersect_warn(self, kwargs)
        self._intersect_warn(other, kwargs)
        dict.update(self, other, **kwargs)

    @staticmethod
    def _intersect_warn(d1, d2):
        intersect = set(d1.keys()) & set(d2.keys())
        if intersect:
            print("Warning: name collision for:",
                  ",".join("%s (%s vs %s)" % (k, type(d1[k]), d2[k])
                           for k in sorted(intersect))) def split_signature(signature):
    signature = signature.strip()
    assert signature.count("(") == 1
    assert signature.count(")") == 1
    assert signature[-1] == ')'
    pos = signature.find('(')
    name = signature[:pos]
    args = signature[pos + 1:-1]
    return name, args def argspec(*args, **kwonlyargs):
    """
    >>> argspec('a, b, c=1, *d, **e, f=None')
    ... # doctest: +NORMALIZE_WHITESPACE
    FullArgSpec(args=['a', 'b', 'c'], varargs='d', varkw='e', defaults=(1,),
                kwonlyargs=['f'], kwonlydefaults={'f': None}, annotations={})
    >>> a = argspec('*, f=None', g=None)
    >>> a.kwonlyargs
    ['f', 'g']
    >>> sorted(a.kwonlydefaults.items())
    [('f', None), ('g', None)]
    >>> argspec('a, *, **b, c=None')
    ... # doctest: +NORMALIZE_WHITESPACE
    FullArgSpec(args=['a'], varargs=None, varkw='b', defaults=None,
                kwonlyargs=['c'], kwonlydefaults={'c': None}, annotations={})
    >>> argspec('a, b, c=1, d=None')
    ... # doctest: +NORMALIZE_WHITESPACE
    FullArgSpec(args=['a', 'b', 'c', 'd'], varargs=None, varkw=None,
                defaults=(1, None), kwonlyargs=[], kwonlydefaults={},
                annotations={})
    >>> argspec('')
    ... # doctest: +NORMALIZE_WHITESPACE
    FullArgSpec(args=[], varargs=None, varkw=None, defaults=None, kwonlyargs=[],
                kwonlydefaults={}, annotations={})
    >>> str(argspec('a, b, c=1, *d, **e, f=None'))
    'a, b, c=1, *d, **e, f=None'
    """
    if len(args) == 1 and isinstance(args[0], basestring):
        str_args = args[0].split(',') if args[0] else []
        str_args = [a.strip().split('=') for a in str_args]
        args = [(a[0], ast.literal_eval(a[1])) if len(a) > 1 else a[0]
                for a in str_args]

        def star(a):
            return isinstance(a, basestring) and '*' in a
        if any(star(a) for a in args):
            while not star(args[-1]):
                k, v = args.pop()
                assert k not in kwonlyargs, "several kwonlyargs named %s" % k
                kwonlyargs[k] = v
    return _argspec(*args, **kwonlyargs) class UserDeprecationWarning(UserWarning):
    pass class ValueForPeriod(TimeFunction):
    funcname = 'value_for_period'

    def compute(self, context, expr, period, missing='auto'):
        return self.value_for_period(expr, period, context, missing)

    dtype = firstarg_dtype

import collections
import re
import sys
import warnings
import numpy as np
import tables
from liam2 import config
from liam2.compat import basestring
from liam2.data import merge_arrays, get_fields, ColumnArray, index_table, build_period_array
from liam2.expr import (Variable, VariableMethodHybrid, GlobalVariable, GlobalTable, GlobalArray, Expr, BinaryOp,
                        MethodSymbol, normalize_type)
from liam2.exprtools import parse
from liam2.process import Assignment, ProcessGroup, While, Function, Return
from liam2.utils import (count_occurrences, field_str_to_type, size2str, WarnOverrideDict, split_signature, argspec,
                         UserDeprecationWarning)
from liam2.tfunc import ValueForPeriod
        from liam2.links import Many2One, One2Many
        from liam2.links import PrefixingLink
        from liam2.tfunc import Lag
        from liam2.links import LinkGet

        if macro_strings is None:
            macro_strings = {}
        self.macro_strings = macro_strings

        self.process_strings = process_strings
        self.processes = None

        self.expectedrows = tables.parameters.EXPECTED_ROWS_TABLE
        self.table = None
        self.input_table = None

        self.indexed_input_table = None
        self.indexed_output_table = None

        self.input_rows = {}
        # TODO: it is unnecessary to keep periods which have already been
        # simulated, because (currently) when we go back in time, we always go
        # back using the output table... but periods before the start_period
        # are only present in input_index
        self.input_index = {}

        self.output_rows = {}
        self.output_index = {}
        self.output_index_node = None

        self.base_period = None
        # we need a separate field, instead of using array['period'] to be able
        # to get the period even when the array is empty.
        self.array_period = array_period
        self.array = array

        self.lag_fields = []
        self.array_lag = None

        self.num_tmp = 0
        self.temp_variables = {}
        self.id_to_rownum = None
        if array is not None:
            rows_per_period, index_per_period = index_table(array)
            self.input_rows = rows_per_period
            self.output_rows = rows_per_period
            self.input_index = index_per_period
            self.output_index = index_per_period
            self.id_to_rownum = index_per_period[array_period]
        self._variables = None
        self._methods = None

    @classmethod
    def from_yaml(cls, ent_name, entity_def):

        # YAML "ordered dict" syntax returns a list of dict and we want a list
        # of tuples
        # FIXME: if "fields" key is present but no field is defined,
        # entity_def.get('fields', []) returns None and this breaks
        fields_def = [list(d.items())[0] for d in entity_def.get('fields', [])]

        link_defs = entity_def.get('links', {})
        str2class = {'one2many': One2Many, 'many2one': Many2One}
        links = dict((name,
                      str2class[l['type']](name, l['field'], l['target']))
                     for name, l in link_defs.items())

        return Entity(ent_name, fields_def, links,
                      entity_def.get('macros', {}),
                      entity_def.get('processes', {}))

    # noinspection PyProtectedMember
    def attach_and_resolve_links(self, entities):
        for link in self.links.values():
            link._attach(self)
            link._resolve_target(entities)

    @property
    def local_var_names(self):
        return set(self.temp_variables.keys()) - set(self.variables.keys())

    @classmethod
    def from_table(cls, table):
        return Entity(table.name, get_fields(table), links={}, macro_strings={},
                      process_strings={})

    # TODO: in_process_group argument will become useless if we drop support for detecting processes defined
    # outside functions
    @staticmethod
    def is_expr(k, v, in_process_group=False):
        if in_process_group:
            # I prefer listing bool explicitly even if not necessary because isinstance(True, int) is True
            iswhile = k is not None and re.match("while[ (].*", k)
            return not iswhile and isinstance(v, (basestring, bool, int, float, list, dict))
        else:
            return isinstance(v, (basestring, bool, int, float))

    @staticmethod
    def collect_predictors(items, in_process_group=False):
        # this excludes lists (functions) and dict (while, ...)
        return [k for k, v in items if k is not None and Entity.is_expr(k, v, in_process_group)]

    @property
    def variables(self):
        if self._variables is None:
            if self.process_strings:
                processes = list(self.process_strings.items())
            else:
                processes = []

            # names of all processes (hybrid or not) of the entity
            process_names = set(k for k, v in processes if k is not None)

            # names of all entity variables (temporary or not) which are set
            # globally
            all_entity_variables = set(self.collect_predictors(processes))

            field_names = set(self.fields.names)

            # normal fields (non-callable/no hybrid variable-function for them)
            variables = dict((name, Variable(self, name, type_))
                             for name, type_ in self.fields.name_types
                             if name in field_names - process_names)


            if config.debug: