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

'''
.. deprecated:: 0.6
    Use :mod:`mosql.query` instead.

.. warning::
    This module will be removed in version 1.0.

It contains the common SQL builders.

.. versionchanged:: 0.2
    It is renamed from ``common``.

.. versionchanged:: 0.1.6
    It is rewritten for using new :mod:`mosql.util`, but it is compatible with
    old version.

.. autosummary ::
    select
    insert
    update
    delete
    join
    or_

It is designed for building the standard SQL statement and tested in PostgreSQL.

.. note::
    If you use MySQL, here is a patch for MySQL --- :mod:`mosql.mysql`.
'''

# --- the removal warning ---
from .util import warning
warning('mosql.build will be removed in version 1.0.')
# --- end ---

__all__ = ['select', 'insert', 'delete', 'update', 'join', 'or_']

from .util import *

# defines formatting chains
single_value      = (value, )
single_identifier = (identifier, )
single_identifier_as = (identifier_as, )
identifier_list   = (identifier, concat_by_comma)
identifier_as_list   = (identifier_as, concat_by_comma)
identifier_dir_list   = (identifier_dir, concat_by_comma)
column_list       = (identifier, concat_by_comma, paren)
value_list        = (value, concat_by_comma, paren)
where_list        = (build_where, )
set_list          = (build_set, )
statement_list    = (concat_by_space, )

# insert

insert    = Clause('insert into', single_identifier)
columns   = Clause('columns'    , column_list, hidden=True)
values    = Clause('values'     , value_list)
returning = Clause('returning'  , identifier_as_list)
on_duplicate_key_update = Clause('on duplicate key update', set_list)

insert_into_stat = Statement([insert, columns, values, returning, on_duplicate_key_update])

def insert(table, set=None, values=None, **clauses_args):
    '''It generates the SQL statement, ``insert into ...``.

    The following usages generate the same SQL statement:

    >>> print insert('person', {'person_id': 'mosky', 'name': 'Mosky Liu'})
    INSERT INTO "person" ("person_id", "name") VALUES ('mosky', 'Mosky Liu')

    >>> print insert('person', (('person_id', 'mosky'), ('name', 'Mosky Liu')))
    INSERT INTO "person" ("person_id", "name") VALUES ('mosky', 'Mosky Liu')

    >>> print insert('person', ('person_id', 'name'), ('mosky', 'Mosky Liu'))
    INSERT INTO "person" ("person_id", "name") VALUES ('mosky', 'Mosky Liu')

    The columns is ignorable:

    >>> print insert('person', values=('mosky', 'Mosky Liu'))
    INSERT INTO "person" VALUES ('mosky', 'Mosky Liu')

    The :func:`insert`, :func:`update` and :func:`delete` support ``returning``.

    >>> print insert('person', {'person_id': 'mosky', 'name': 'Mosky Liu'}, returning=raw('*'))
    INSERT INTO "person" ("person_id", "name") VALUES ('mosky', 'Mosky Liu') RETURNING *

    The MySQL-specific "on duplicate key update" is also supported:

    >>> print insert('person', values=('mosky', 'Mosky Liu'), on_duplicate_key_update={'name': 'Mosky Liu'})
    INSERT INTO "person" VALUES ('mosky', 'Mosky Liu') ON DUPLICATE KEY UPDATE "name"='Mosky Liu'
    '''

    clauses_args['insert into'] = table

    if values is None:
        if hasattr(set, 'items'):
            pairs = set.items()
        else:
            pairs = set
        clauses_args['columns'], clauses_args['values'] = zip(*pairs)
    else:
        clauses_args['columns'] = set
        clauses_args['values']  = values

    if 'on_duplicate_key_update' in clauses_args:
        clauses_args['on duplicate key update'] = clauses_args['on_duplicate_key_update']
        del clauses_args['on_duplicate_key_update']

    return insert_into_stat.format(clauses_args)

# select

select   = Clause('select'  , identifier_as_list)
from_    = Clause('from'    , identifier_as_list)
joins    = Clause('joins'   , statement_list, hidden=True)
where    = Clause('where'   , where_list)
group_by = Clause('group by', identifier_list)
having   = Clause('having'  , where_list)
order_by = Clause('order by', identifier_dir_list)
limit    = Clause('limit'   , single_value)
offset   = Clause('offset'  , single_value)

select_stat = Statement([select, from_, joins, where, group_by, having, order_by, limit, offset])

def select(table, where=None, select=None, **clauses_args):
    '''It generates the SQL statement, ``select ...`` .

    .. versionchanged:: 0.1.6
        The clause argument, ``join``, is renamed to ``joins``.

    The following usages generate the same SQL statement.

    >>> print select('person', {'person_id': 'mosky'})
    SELECT * FROM "person" WHERE "person_id" = 'mosky'

    >>> print select('person', (('person_id', 'mosky'), ))
    SELECT * FROM "person" WHERE "person_id" = 'mosky'

    It detects the dot in an identifier:

    >>> print select('person', select=('person.person_id', 'person.name'))
    SELECT "person"."person_id", "person"."name" FROM "person"

    Building prepare statement with :class:`mosql.util.param`:

    >>> print select('table', {'custom_param': param('my_param'), 'auto_param': param, 'using_alias': ___})
    SELECT * FROM "table" WHERE "auto_param" = %(auto_param)s AND "using_alias" = %(using_alias)s AND "custom_param" = %(my_param)s

    You can also specify the ``group_by``, ``having``, ``order_by``, ``limit``
    and ``offset`` in the keyword arguments. Here are some examples:

    >>> print select('person', {'name like': 'Mosky%'}, group_by=('age', ))
    SELECT * FROM "person" WHERE "name" LIKE 'Mosky%' GROUP BY "age"

    >>> print select('person', {'name like': 'Mosky%'}, order_by=('age', ))
    SELECT * FROM "person" WHERE "name" LIKE 'Mosky%' ORDER BY "age"

    >>> print select('person', {'name like': 'Mosky%'}, order_by=('age desc', ))
    SELECT * FROM "person" WHERE "name" LIKE 'Mosky%' ORDER BY "age" DESC

    >>> print select('person', {'name like': 'Mosky%'}, order_by=('age ; DROP person; --', ))
    Traceback (most recent call last):
        ...
    OptionError: this option is not allowed: '; DROP PERSON; --'

    .. seealso ::
        The options allowed --- :attr:`mosql.util.allowed_options`.

    >>> print select('person', {'name like': 'Mosky%'}, limit=3, offset=1)
    SELECT * FROM "person" WHERE "name" LIKE 'Mosky%' LIMIT 3 OFFSET 1

    The operators are also supported:

    >>> print select('person', {'person_id': ('andy', 'bob')})
    SELECT * FROM "person" WHERE "person_id" IN ('andy', 'bob')

    >>> print select('person', {'name': None})
    SELECT * FROM "person" WHERE "name" IS NULL

    >>> print select('person', {'name like': 'Mosky%', 'age >': 20})
    SELECT * FROM "person" WHERE "age" > 20 AND "name" LIKE 'Mosky%'

    >>> print select('person', {"person_id = '' OR true; --": 'mosky'})
    Traceback (most recent call last):
        ...
    OperatorError: this operator is not allowed: "= '' OR TRUE; --"

    .. seealso ::
        The operators allowed --- :attr:`mosql.util.allowed_operators`.

    If you want to use functions, wrap it with :class:`mosql.util.raw`:

    >>> print select('person', select=raw('count(*)'), group_by=('age', ))
    SELECT count(*) FROM "person" GROUP BY "age"

    .. warning ::
        It is your responsibility to ensure that your SQL queries are properly escaped if you use :class:`mosql.util.raw`.

    .. seealso ::
        How it builds the where clause --- :func:`mosql.util.build_where`
    '''

    clauses_args['from']   = table
    clauses_args['where']  = where
    clauses_args['select'] = star if select is None else select

    if 'order_by' in clauses_args:
        clauses_args['order by'] = clauses_args['order_by']
        del clauses_args['order_by']

    if 'group_by' in clauses_args:
        clauses_args['group by'] = clauses_args['group_by']
        del clauses_args['group_by']

    return select_stat.format(clauses_args)

# update

update = Clause('update', single_identifier_as)
set    = Clause('set'   , set_list)

update_stat = Statement([update, set, where, returning])

def update(table, where=None, set=None, **clauses_args):
    '''It generates the SQL statement, ``update ...`` .

    The following usages generate the same SQL statement.

    >>> print update('person', {'person_id': 'mosky'}, {'name': 'Mosky Liu'})
    UPDATE "person" SET "name"='Mosky Liu' WHERE "person_id" = 'mosky'

    >>> print update('person', (('person_id', 'mosky'), ), (('name', 'Mosky Liu'),) )
    UPDATE "person" SET "name"='Mosky Liu' WHERE "person_id" = 'mosky'

    .. seealso ::
        How it builds the where clause --- :func:`mosql.util.build_set`
    '''

    clauses_args['update'] = table
    clauses_args['where']  = where
    clauses_args['set']    = set

    return update_stat.format(clauses_args)

# delete from

delete = Clause('delete from', single_identifier_as)

delete_stat = Statement([delete, where, returning])

def delete(table, where=None, **clauses_args):
    '''It generates the SQL statement, ``delete from ...`` .

    The following usages generate the same SQL statement.

    >>> print delete('person', {'person_id': 'mosky'})
    DELETE FROM "person" WHERE "person_id" = 'mosky'

    >>> print delete('person', (('person_id', 'mosky'), ))
    DELETE FROM "person" WHERE "person_id" = 'mosky'
    '''

    clauses_args['delete from'] = table
    clauses_args['where'] = where

    return delete_stat.format(clauses_args)

# join

join  = Clause('join' , single_identifier_as)
type  = Clause('type' , tuple(), hidden=True)
on    = Clause('on'   , (build_on, ))
using = Clause('using', column_list)

join_stat = Statement([type, join, on, using])

def join(table, using=None, on=None, type=None, **clauses_args):
    '''It generates the SQL statement, ``... join ...`` .

    .. versionadded :: 0.1.6

    >>> print select('person', joins=join('detail'))
    SELECT * FROM "person" NATURAL JOIN "detail"

    >>> print select('person', joins=join('detail', using=('person_id', )))
    SELECT * FROM "person" INNER JOIN "detail" USING ("person_id")

    >>> print select('person', joins=join('detail', on={'person.person_id': 'detail.person_id'}))
    SELECT * FROM "person" INNER JOIN "detail" ON "person"."person_id" = "detail"."person_id"

    >>> print select('person', joins=join('detail', type='cross'))
    SELECT * FROM "person" CROSS JOIN "detail"

    .. seealso ::
        How it builds the where clause --- :func:`mosql.util.build_on`
    '''

    clauses_args['join'] = table
    clauses_args['using'] = using
    clauses_args['on'] = on

    if not type:
        if using or on:
            clauses_args['type'] = 'INNER'
        else:
            clauses_args['type'] = 'NATURAL'
    else:
        clauses_args['type'] = type.upper()

    return join_stat.format(clauses_args)

# or

def or_(*conditions):
    '''It concats the conditions by ``OR``.

    .. versionadded :: 0.1.6

    >>> print or_({'person_id': 'andy'}, {'person_id': 'bob'})
    "person_id" = 'andy' OR "person_id" = 'bob'
    '''

    return concat_by_or(build_where(c) for c in conditions)

if __name__ == '__main__':
    import doctest
    doctest.testmod()

    # benchmark

    #from timeit import timeit
    #from functools import partial
    #timeit = partial(timeit, number=100000)

    #import mosql.util

    #print timeit(lambda: select('person', {'name': 'Mosky Liu'}, ('person_id', 'name'), limit=10, order_by='person_id'))
    ## -> 4.97957897186

    #print timeit(lambda: select('person', {'name': 'Mosky Liu'}, ('person.person_id', 'person.name'), limit=10, order_by='person_id'))
    ## -> 5.33279800415

    #mosql.util.delimit_identifier = None
    #print timeit(lambda: select('person', {'name': 'Mosky Liu'}, ('person_id', 'name'), limit=10, order_by='person_id'))
    ## -> 3.94950485229

    ##from mosql.common import select as old_select

    ##print timeit(lambda: old_select('person', {'name': 'Mosky Liu'}, ('person_id', 'name'), limit=10, order_by='person_id'))
    ### -> 6.79131507874
<commit_msg>align<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

'''
.. deprecated:: 0.6
    Use :mod:`mosql.query` instead.

.. warning::
    This module will be removed in version 1.0.

It contains the common SQL builders.

.. versionchanged:: 0.2
    It is renamed from ``common``.

.. versionchanged:: 0.1.6
    It is rewritten for using new :mod:`mosql.util`, but it is compatible with
    old version.

.. autosummary ::
    select
    insert
    update
    delete
    join
    or_

It is designed for building the standard SQL statement and tested in PostgreSQL.

.. note::
    If you use MySQL, here is a patch for MySQL --- :mod:`mosql.mysql`.
'''

# --- the removal warning ---
from .util import warning
warning('mosql.build will be removed in version 1.0.')
# --- end ---

__all__ = ['select', 'insert', 'delete', 'update', 'join', 'or_']

from .util import *

# defines formatting chains
single_value          = (value, )
single_identifier     = (identifier, )
single_identifier_as  = (identifier_as, )
identifier_list       = (identifier, concat_by_comma)
identifier_as_list    = (identifier_as, concat_by_comma)
identifier_dir_list   = (identifier_dir, concat_by_comma)
column_list           = (identifier, concat_by_comma, paren)
value_list            = (value, concat_by_comma, paren)
where_list            = (build_where, )
set_list              = (build_set, )
statement_list        = (concat_by_space, )

# insert

insert    = Clause('insert into', single_identifier)
columns   = Clause('columns'    , column_list, hidden=True)
values    = Clause('values'     , value_list)
returning = Clause('returning'  , identifier_as_list)
on_duplicate_key_update = Clause('on duplicate key update', set_list)

insert_into_stat = Statement([insert, columns, values, returning, on_duplicate_key_update])

def insert(table, set=None, values=None, **clauses_args):
    '''It generates the SQL statement, ``insert into ...``.

    The following usages generate the same SQL statement:

    >>> print insert('person', {'person_id': 'mosky', 'name': 'Mosky Liu'})
    INSERT INTO "person" ("person_id", "name") VALUES ('mosky', 'Mosky Liu')

    >>> print insert('person', (('person_id', 'mosky'), ('name', 'Mosky Liu')))
    INSERT INTO "person" ("person_id", "name") VALUES ('mosky', 'Mosky Liu')

    >>> print insert('person', ('person_id', 'name'), ('mosky', 'Mosky Liu'))
    INSERT INTO "person" ("person_id", "name") VALUES ('mosky', 'Mosky Liu')

    The columns is ignorable:

    >>> print insert('person', values=('mosky', 'Mosky Liu'))
    INSERT INTO "person" VALUES ('mosky', 'Mosky Liu')

    The :func:`insert`, :func:`update` and :func:`delete` support ``returning``.

    >>> print insert('person', {'person_id': 'mosky', 'name': 'Mosky Liu'}, returning=raw('*'))
    INSERT INTO "person" ("person_id", "name") VALUES ('mosky', 'Mosky Liu') RETURNING *

    The MySQL-specific "on duplicate key update" is also supported:

    >>> print insert('person', values=('mosky', 'Mosky Liu'), on_duplicate_key_update={'name': 'Mosky Liu'})
    INSERT INTO "person" VALUES ('mosky', 'Mosky Liu') ON DUPLICATE KEY UPDATE "name"='Mosky Liu'
    '''

    clauses_args['insert into'] = table

    if values is None:
        if hasattr(set, 'items'):
            pairs = set.items()
        else:
            pairs = set
        clauses_args['columns'], clauses_args['values'] = zip(*pairs)
    else:
        clauses_args['columns'] = set
        clauses_args['values']  = values

    if 'on_duplicate_key_update' in clauses_args:
        clauses_args['on duplicate key update'] = clauses_args['on_duplicate_key_update']
        del clauses_args['on_duplicate_key_update']

    return insert_into_stat.format(clauses_args)

# select

select   = Clause('select'  , identifier_as_list)
from_    = Clause('from'    , identifier_as_list)
joins    = Clause('joins'   , statement_list, hidden=True)
where    = Clause('where'   , where_list)
group_by = Clause('group by', identifier_list)
having   = Clause('having'  , where_list)
order_by = Clause('order by', identifier_dir_list)
limit    = Clause('limit'   , single_value)
offset   = Clause('offset'  , single_value)

select_stat = Statement([select, from_, joins, where, group_by, having, order_by, limit, offset])

def select(table, where=None, select=None, **clauses_args):
    '''It generates the SQL statement, ``select ...`` .

    .. versionchanged:: 0.1.6
        The clause argument, ``join``, is renamed to ``joins``.

    The following usages generate the same SQL statement.

    >>> print select('person', {'person_id': 'mosky'})
    SELECT * FROM "person" WHERE "person_id" = 'mosky'

    >>> print select('person', (('person_id', 'mosky'), ))
    SELECT * FROM "person" WHERE "person_id" = 'mosky'

    It detects the dot in an identifier:

    >>> print select('person', select=('person.person_id', 'person.name'))
    SELECT "person"."person_id", "person"."name" FROM "person"

    Building prepare statement with :class:`mosql.util.param`:

    >>> print select('table', {'custom_param': param('my_param'), 'auto_param': param, 'using_alias': ___})
    SELECT * FROM "table" WHERE "auto_param" = %(auto_param)s AND "using_alias" = %(using_alias)s AND "custom_param" = %(my_param)s

    You can also specify the ``group_by``, ``having``, ``order_by``, ``limit``
    and ``offset`` in the keyword arguments. Here are some examples:

    >>> print select('person', {'name like': 'Mosky%'}, group_by=('age', ))
    SELECT * FROM "person" WHERE "name" LIKE 'Mosky%' GROUP BY "age"

    >>> print select('person', {'name like': 'Mosky%'}, order_by=('age', ))
    SELECT * FROM "person" WHERE "name" LIKE 'Mosky%' ORDER BY "age"

    >>> print select('person', {'name like': 'Mosky%'}, order_by=('age desc', ))
    SELECT * FROM "person" WHERE "name" LIKE 'Mosky%' ORDER BY "age" DESC

    >>> print select('person', {'name like': 'Mosky%'}, order_by=('age ; DROP person; --', ))
    Traceback (most recent call last):
        ...
    OptionError: this option is not allowed: '; DROP PERSON; --'

    .. seealso ::
        The options allowed --- :attr:`mosql.util.allowed_options`.

    >>> print select('person', {'name like': 'Mosky%'}, limit=3, offset=1)
    SELECT * FROM "person" WHERE "name" LIKE 'Mosky%' LIMIT 3 OFFSET 1

    The operators are also supported:

    >>> print select('person', {'person_id': ('andy', 'bob')})
    SELECT * FROM "person" WHERE "person_id" IN ('andy', 'bob')

    >>> print select('person', {'name': None})
    SELECT * FROM "person" WHERE "name" IS NULL

    >>> print select('person', {'name like': 'Mosky%', 'age >': 20})
    SELECT * FROM "person" WHERE "age" > 20 AND "name" LIKE 'Mosky%'

    >>> print select('person', {"person_id = '' OR true; --": 'mosky'})
    Traceback (most recent call last):
        ...
    OperatorError: this operator is not allowed: "= '' OR TRUE; --"

    .. seealso ::
        The operators allowed --- :attr:`mosql.util.allowed_operators`.

    If you want to use functions, wrap it with :class:`mosql.util.raw`:

    >>> print select('person', select=raw('count(*)'), group_by=('age', ))
    SELECT count(*) FROM "person" GROUP BY "age"

    .. warning ::
        It is your responsibility to ensure that your SQL queries are properly escaped if you use :class:`mosql.util.raw`.

    .. seealso ::
        How it builds the where clause --- :func:`mosql.util.build_where`
    '''

    clauses_args['from']   = table
    clauses_args['where']  = where
    clauses_args['select'] = star if select is None else select

    if 'order_by' in clauses_args:
        clauses_args['order by'] = clauses_args['order_by']
        del clauses_args['order_by']

    if 'group_by' in clauses_args:
        clauses_args['group by'] = clauses_args['group_by']
        del clauses_args['group_by']

    return select_stat.format(clauses_args)

# update

update = Clause('update', single_identifier_as)
set    = Clause('set'   , set_list)

update_stat = Statement([update, set, where, returning])

def update(table, where=None, set=None, **clauses_args):
    '''It generates the SQL statement, ``update ...`` .

    The following usages generate the same SQL statement.

    >>> print update('person', {'person_id': 'mosky'}, {'name': 'Mosky Liu'})
    UPDATE "person" SET "name"='Mosky Liu' WHERE "person_id" = 'mosky'

    >>> print update('person', (('person_id', 'mosky'), ), (('name', 'Mosky Liu'),) )
    UPDATE "person" SET "name"='Mosky Liu' WHERE "person_id" = 'mosky'

    .. seealso ::
        How it builds the where clause --- :func:`mosql.util.build_set`
    '''

    clauses_args['update'] = table
    clauses_args['where']  = where
    clauses_args['set']    = set

    return update_stat.format(clauses_args)

# delete from

delete = Clause('delete from', single_identifier_as)

delete_stat = Statement([delete, where, returning])

def delete(table, where=None, **clauses_args):
    '''It generates the SQL statement, ``delete from ...`` .

    The following usages generate the same SQL statement.

    >>> print delete('person', {'person_id': 'mosky'})
    DELETE FROM "person" WHERE "person_id" = 'mosky'

    >>> print delete('person', (('person_id', 'mosky'), ))
    DELETE FROM "person" WHERE "person_id" = 'mosky'
    '''

    clauses_args['delete from'] = table
    clauses_args['where'] = where

    return delete_stat.format(clauses_args)

# join

join  = Clause('join' , single_identifier_as)
type  = Clause('type' , tuple(), hidden=True)
on    = Clause('on'   , (build_on, ))
using = Clause('using', column_list)

join_stat = Statement([type, join, on, using])

def join(table, using=None, on=None, type=None, **clauses_args):
    '''It generates the SQL statement, ``... join ...`` .

    .. versionadded :: 0.1.6

    >>> print select('person', joins=join('detail'))
    SELECT * FROM "person" NATURAL JOIN "detail"

    >>> print select('person', joins=join('detail', using=('person_id', )))
    SELECT * FROM "person" INNER JOIN "detail" USING ("person_id")

    >>> print select('person', joins=join('detail', on={'person.person_id': 'detail.person_id'}))
    SELECT * FROM "person" INNER JOIN "detail" ON "person"."person_id" = "detail"."person_id"

    >>> print select('person', joins=join('detail', type='cross'))
    SELECT * FROM "person" CROSS JOIN "detail"

    .. seealso ::
        How it builds the where clause --- :func:`mosql.util.build_on`
    '''

    clauses_args['join'] = table
    clauses_args['using'] = using
    clauses_args['on'] = on

    if not type:
        if using or on:
            clauses_args['type'] = 'INNER'
        else:
            clauses_args['type'] = 'NATURAL'
    else:
        clauses_args['type'] = type.upper()

    return join_stat.format(clauses_args)

# or

def or_(*conditions):
    '''It concats the conditions by ``OR``.

    .. versionadded :: 0.1.6

    >>> print or_({'person_id': 'andy'}, {'person_id': 'bob'})
    "person_id" = 'andy' OR "person_id" = 'bob'
    '''

    return concat_by_or(build_where(c) for c in conditions)

if __name__ == '__main__':
    import doctest
    doctest.testmod()

    # benchmark

    #from timeit import timeit
    #from functools import partial
    #timeit = partial(timeit, number=100000)

    #import mosql.util

    #print timeit(lambda: select('person', {'name': 'Mosky Liu'}, ('person_id', 'name'), limit=10, order_by='person_id'))
    ## -> 4.97957897186

    #print timeit(lambda: select('person', {'name': 'Mosky Liu'}, ('person.person_id', 'person.name'), limit=10, order_by='person_id'))
    ## -> 5.33279800415

    #mosql.util.delimit_identifier = None
    #print timeit(lambda: select('person', {'name': 'Mosky Liu'}, ('person_id', 'name'), limit=10, order_by='person_id'))
    ## -> 3.94950485229

    ##from mosql.common import select as old_select

    ##print timeit(lambda: old_select('person', {'name': 'Mosky Liu'}, ('person_id', 'name'), limit=10, order_by='person_id'))
    ### -> 6.79131507874
<|endoftext|>


<commit_before>"""
Test module for queries on datasets
===================================

:Author:   Ivan Vilata i Balaguer
:Contact:  ivilata@carabos.com
:Created:  2006-10-19
:License:  BSD
:Revision: $Id$
"""

import re
import new
import unittest

import numpy

import tables
import tables.tests.common as tests

row_period = 50
"""Maximum number of unique rows before they start cycling."""
md_shape = (2, 2)
"""Shape of multidimensional fields."""

_maxnvalue = row_period + numpy.product(md_shape) - 1
_strlen = int(numpy.log10(_maxnvalue-1)) + 1

str_format = '%%0%dd' % _strlen
"""Format of string values."""

ptype_info = {
    'Bool': (numpy.bool_, bool),
    'Int8': (numpy.int8, int), 'UInt8': (numpy.uint8, int),
    'Int16': (numpy.int16, int), 'UInt16': (numpy.uint16, int),
    'Int32': (numpy.int32, int), 'UInt32': (numpy.uint32, long),
    'Int64': (numpy.int64, long), 'UInt64': (numpy.uint64, long),
    'Float32': (numpy.float32, float), 'Float64': (numpy.float32, float),
    'Complex32': (numpy.complex64, complex),
    'Complex64': (numpy.complex128, complex),
    'Time32': (numpy.int32, int), 'Time64': (numpy.float64, float),
    'Enum': (numpy.uint8, int),  # just for these tests
    'String': (numpy.dtype('S%s' % _strlen), str) }  # just for these tests
"""NumPy and Numexpr type for each PyTables type that will be tested."""

dtype_from_ptype = dict(
    (dtype, info[0]) for (dtype, info) in ptype_info.iteritems() )
"""Maps PyTables types to NumPy data types."""
nxtype_from_ptype = dict(
    (dtype, info[1]) for (dtype, info) in ptype_info.iteritems() )
"""Maps PyTables types to Numexpr data types."""

heavy_types = ['UInt8', 'Int16', 'UInt16', 'Float32', 'Complex32']
"""Python types to be tested only in heavy mode."""

if not tests.heavy:
    for ptype in heavy_types:
        for tdict in ptype_info, dtype_from_ptype, nxtype_from_ptype:
            del tdict[ptype]

enum = tables.Enum(dict(('n%d' % i, i) for i in range(_maxnvalue)))
"""Enumerated type to be used in tests."""

def append_columns(classdict, shape=()):
    """
    Append a ``Col`` of each PyTables data type to the `classdict`.

    A column of a certain TYPE gets called ``cTYPE``.  The number of
    added columns is returned.
    """
    for (itype, ptype) in enumerate(sorted(ptype_info.iterkeys())):
        colpos = itype + 1
        colname = 'c%s' % ptype
        colclass = getattr(tables, '%sCol' % ptype)
        colargs, colkwargs = [], {'shape': shape, 'pos': colpos}
        if ptype == 'Enum':
            colargs = [enum, enum(0), dtype_from_ptype[ptype]] + colargs
        elif ptype == 'String':
            colkwargs.update({'length': dtype_from_ptype[ptype].itemsize})
        classdict[colname] = colclass(*colargs, **colkwargs)
    ncols = colpos
    return ncols

def nested_description(classname, pos, shape=()):
    """
    Return a nested column description with all PyTables data types.

    A column of a certain TYPE gets called ``cTYPE``.  The nested
    column will be placed in the position indicated by `pos`.
    """
    classdict = {}
    append_columns(classdict, shape=shape)
    classdict['_v_pos'] = pos
    return new.classobj(classname, (tables.IsDescription,), classdict)

def table_description(classname, nclassname, shape=()):
    """
    Return a table description for testing queries.

    The description consists of all PyTables data types, both in the
    top level and in the ``cNested`` nested column.  A column of a
    certain TYPE gets called ``cTYPE``.  An extra integer column
    ``cExtra`` is also provided.  If a `shape` is given, it will be
    used for all columns.
    """
    classdict = {}
    colpos = append_columns(classdict, shape)

    ndescr = nested_description(nclassname, colpos, shape=shape)
    classdict['cNested'] = ndescr
    colpos += 1

    extracol = tables.IntCol(shape=shape, pos=colpos)
    classdict['cExtra'] = extracol
    colpos += 1

    return new.classobj(classname, (tables.IsDescription,), classdict)

TableDescription = table_description(
    'TableDescription', 'NestedDescription' )
"""Unidimensional table description for testing queries."""

MDTableDescription = table_description(
    'MDTableDescription', 'MDNestedDescription', shape=md_shape )
"""Multidimensional table description for testing queries."""

table_data = {}
"""Cached table data for a given shape and number of rows."""
# Data is cached because computing it row by row is quite slow.  Hop!

def fill_table(table, shape, nrows):
    """
    Fill the given `table` with `nrows` rows of data.

    Values in the i-th row (where 0 <= i < `row_period`) for a
    multidimensional field with M elements span from i to i+M-1.  For
    subsequent rows, values repeat cyclically.

    The same goes for the ``cExtra`` column, but values range from
    -`row_period`/2 to +`row_period`/2.
    """
    # Reuse already computed data if possible.
    tdata = table_data.get((shape, nrows))
    if tdata is not None:
        table.append(tdata)
        table.flush()
        return

    size = int(numpy.product(shape))

    row, value = table.row, 0
    for nrow in xrange(nrows):
        data = numpy.arange(value, value + size).reshape(shape)
        for (ptype, dtype) in dtype_from_ptype.iteritems():
            colname = 'c%s' % ptype
            ncolname = 'cNested/%s' % colname
            if ptype == 'Bool':
                coldata = data > (row_period / 2)
            elif ptype == 'String':
                sdata = [str_format % x for x in range(value, value + size)]
                coldata = numpy.array(sdata, dtype=dtype).reshape(shape)
            else:
                coldata = numpy.asarray(data, dtype=dtype)
            row[ncolname] = row[colname] = coldata
            row['cExtra'] = data - (row_period / 2)
        row.append()
        value += 1
        if value == row_period:
            value = 0
    table.flush()

    # Make computed data reusable.
    tdata = table.read()
    table_data[(shape, nrows)] = tdata

class TableQueryTestCase(tests.TempFileMixin, tests.PyTablesTestCase):

    """
    Test querying table data.

    Sub-classes must define the following attributes:

    ``tableDescription``
        The description of the table to be created.
    ``shape``
        The shape of data fields in the table.
    ``nrows``
        The number of data rows to be generated for the table.

    Sub-classes may redefine the following attributes:

    ``indexed``
        Whether columns shall be indexed, if possible.  Default is not
        to index them.
    ``optlevel``
        The level of optimisation of column indexes.  Default is 0.
    """

    indexed = False
    optlevel = 0

    colNotIndexable_re = re.compile(r"\bcan not be indexed\b")
    condNotBoolean_re = re.compile(r"\bdoes not have a boolean type\b")

    def createIndexes(self, colname, ncolname):
        if not self.indexed:
            return
        try:
            for acolname in [colname, ncolname]:
                acolumn = self.table.colinstances[acolname]
                acolumn.createIndex(optlevel=self.optlevel, testmode=True)
        except TypeError, te:
            if self.colNotIndexable_re.search(str(te)):
                raise tests.SkipTest  # can't be indexed, nothing new to test
            raise
        except NotImplementedError:
            raise tests.SkipTest  # column does not support indexing yet

    def setUp(self):
        super(TableQueryTestCase, self).setUp()
        self.table = table = self.h5file.createTable(
            '/', 'test', self.tableDescription, expectedrows=self.nrows )
        fill_table(table, self.shape, self.nrows)

    ## XXX Need some standard checks on query usage.

operators = [None, '<', '==', '!=']
"""Comparison operators to check with different types."""
if tests.heavy:
    operators += ['~', '<=', '>=', '>']
## XXX Need to check por operator pairs.
left_bound = row_period / 4
"""Operand of left side operator in comparisons with operator pairs."""
right_bound = row_period * 3 / 4
"""Operand of right side operator in comparisons with operator pairs."""
extra_conditions = ['', '& ((cExtra+1) > 0)', '| ((cExtra+1) > 0)']
"""Extra conditions to append to comparison conditions."""

def create_test_method(ptype, op, extracond):
    bound = right_bound
    if ptype == 'String':
        bound = str_format % bound
    bound = nxtype_from_ptype[ptype](bound)
    condvars = {'bound': bound}

    colname = 'c%s' % ptype
    ncolname = 'cNested/%s' % colname
    if not op:
        cond = colname
    elif op == '~':
        cond = '~(%s)' % colname
    elif op == '<':
        cond = '%s %s %s' % (colname, op, bound)  # variable-constant
    else:
        cond = '%s %s bound' % (colname, op)  # variable-variable
    if extracond:
        cond = '(%s) %s' % (cond, extracond)

    def test_method(self):
        # Replace bitwise operators with their logical counterparts.
        pycond = cond
        for (ptop, pyop) in [('&', 'and'), ('|', 'or'), ('~', 'not')]:
            pycond = pycond.replace(ptop, pyop)
        pycond = compile(pycond, '<string>', 'eval')

        table = self.table
        self.createIndexes(colname, ncolname)

        reflen = None
        ## XXX Better try to check retrieved data with ``readWhere()``.
        ## XXX Better yet, try all query methods in parallel!
        # Test that both simple and nested columns work as expected.
        # Because of the way the table is filled, results are the same.
        for acolname in [colname, ncolname]:
            # First the reference Python version.
            pylen, pyvars = 0, condvars.copy()
            for row in table:
                pyvars[colname] = row[acolname]
                pyvars['cExtra'] = row['cExtra']
                try:
                    isvalidrow = eval(pycond, {}, pyvars)
                except TypeError:
                    raise tests.SkipTest  # type doesn't support operation
                if isvalidrow:
                    pylen += 1
            if reflen is None:  # initialise reference length
                reflen = pylen
            self.assertEqual(pylen, reflen)

            # Then the in-kernel or indexed version.
            ptvars = condvars.copy()
            ptvars[colname] = table.colinstances[acolname]
            ptvars['cExtra'] = table.colinstances['cExtra']
            try:
                ptlen = len([r for r in table.where(cond, ptvars)])
            except TypeError, te:
                if self.condNotBoolean_re.search(str(te)):
                    raise tests.SkipTest  # condition is not boolean
                raise
            except NotImplementedError:
                raise tests.SkipTest  # type doesn't support operation
            self.assertEqual(ptlen, reflen)

    test_method.__doc__ = "Testing ``%s``." % cond
    return test_method

# Create individual tests.  You may restrict which tests are generated
# by replacing the sequences in the ``for`` statements.  For instance:
testn = 0
for ptype in ptype_info:  # for ptype in ['String']:
    for op in operators:  # for op in ['!=']:
        for extracond in extra_conditions:  # for extracond in ['']:
            tmethod = create_test_method(ptype, op, extracond)
            tmethod.__name__ = 'test_a%04d' % testn
            ptmethod = tests.pyTablesTest(tmethod)
            imethod = new.instancemethod(ptmethod, None, TableQueryTestCase)
            setattr(TableQueryTestCase, tmethod.__name__, imethod)
            testn += 1

# Base classes for all queries.
class ScalarTableMixin:
    tableDescription = TableDescription
    shape = ()
class MDTableMixin:
    tableDescription = MDTableDescription
    shape = md_shape

# Base classes for non-indexed queries.
NX_BLOCK_SIZE1 = 128  # from ``interpreter.c`` in Numexpr
NX_BLOCK_SIZE2 = 8  # from ``interpreter.c`` in Numexpr

class SmallTableMixin:
    nrows = row_period * 2
    assert NX_BLOCK_SIZE2 < nrows < NX_BLOCK_SIZE1
    assert nrows % NX_BLOCK_SIZE2 != 0  # to have some residual rows
class BigTableMixin:
    nrows = row_period * 3
    assert nrows > NX_BLOCK_SIZE1 + NX_BLOCK_SIZE2
    assert nrows % NX_BLOCK_SIZE1 != 0
    assert nrows % NX_BLOCK_SIZE2 != 0  # to have some residual rows

# Parameters for non-indexed queries.
table_sizes = ['Small']
if tests.heavy:
    table_sizes += ['Big']
table_ndims = ['Scalar']  # to enable multidimensional testing, include 'MD'

# Non-indexed queries: ``[SB][SM]TQTestCase``.
def niclassdata():
    for size in table_sizes:
        for ndim in table_ndims:
            classname = '%s%sTQTestCase' % (size[0], ndim[0])
            cbasenames = ( '%sTableMixin' % size, '%sTableMixin' % ndim,
                           'TableQueryTestCase' )
            yield (classname, cbasenames, {})

# Base classes for indexed queries.
class SmallITableMixin:
    nrows = 50
class MediumITableMixin:
    nrows = 100
class BigITableMixin:
    nrows = 500

# Parameters for indexed queries.
itable_sizes = ['Small']
itable_optvalues = [0, 2, 3]
if tests.heavy:
    itable_sizes += ['Medium', 'Big']
    itable_optvalues += [7, 9]

# Indexed queries: ``[SMB]I[02379]TQTestCase``.
def iclassdata():
    for size in itable_sizes:
        for optlevel in itable_optvalues:
            classname = '%sI%dTQTestCase' % (size[0], optlevel)
            cbasenames = ( '%sITableMixin' % size, 'ScalarTableMixin',
                           'TableQueryTestCase' )
            yield ( classname, cbasenames,
                    {'optlevel': optlevel, 'indexed': True} )

# Create test classes.
for cdatafunc in [niclassdata, iclassdata]:
    for (cname, cbasenames, cdict) in cdatafunc():
        cbases = tuple(eval(cbase) for cbase in cbasenames)
        class_ = new.classobj(cname, cbases, cdict)
        exec '%s = class_' % cname

def suite():
    """Return a test suite consisting of all the test cases in the module."""

    testSuite = unittest.TestSuite()

    niter = 1
    for i in range(niter):
        for cdatafunc in [niclassdata, iclassdata]:
            for cdata in cdatafunc():
                cname = cdata[0]
                testSuite.addTest(unittest.makeSuite(eval(cname)))

    return testSuite

if __name__ == '__main__':
    unittest.main(defaultTest='suite')
<commit_msg>Added verbose output to `test_queries.py`.<commit_after>"""
Test module for queries on datasets
===================================

:Author:   Ivan Vilata i Balaguer
:Contact:  ivilata@carabos.com
:Created:  2006-10-19
:License:  BSD
:Revision: $Id$
"""

import re
import new
import unittest

import numpy

import tables
import tables.tests.common as tests
from tables.tests.common import verbosePrint as vprint

row_period = 50
"""Maximum number of unique rows before they start cycling."""
md_shape = (2, 2)
"""Shape of multidimensional fields."""

_maxnvalue = row_period + numpy.product(md_shape) - 1
_strlen = int(numpy.log10(_maxnvalue-1)) + 1

str_format = '%%0%dd' % _strlen
"""Format of string values."""

ptype_info = {
    'Bool': (numpy.bool_, bool),
    'Int8': (numpy.int8, int), 'UInt8': (numpy.uint8, int),
    'Int16': (numpy.int16, int), 'UInt16': (numpy.uint16, int),
    'Int32': (numpy.int32, int), 'UInt32': (numpy.uint32, long),
    'Int64': (numpy.int64, long), 'UInt64': (numpy.uint64, long),
    'Float32': (numpy.float32, float), 'Float64': (numpy.float32, float),
    'Complex32': (numpy.complex64, complex),
    'Complex64': (numpy.complex128, complex),
    'Time32': (numpy.int32, int), 'Time64': (numpy.float64, float),
    'Enum': (numpy.uint8, int),  # just for these tests
    'String': (numpy.dtype('S%s' % _strlen), str) }  # just for these tests
"""NumPy and Numexpr type for each PyTables type that will be tested."""

dtype_from_ptype = dict(
    (dtype, info[0]) for (dtype, info) in ptype_info.iteritems() )
"""Maps PyTables types to NumPy data types."""
nxtype_from_ptype = dict(
    (dtype, info[1]) for (dtype, info) in ptype_info.iteritems() )
"""Maps PyTables types to Numexpr data types."""

heavy_types = ['UInt8', 'Int16', 'UInt16', 'Float32', 'Complex32']
"""Python types to be tested only in heavy mode."""

if not tests.heavy:
    for ptype in heavy_types:
        for tdict in ptype_info, dtype_from_ptype, nxtype_from_ptype:
            del tdict[ptype]

enum = tables.Enum(dict(('n%d' % i, i) for i in range(_maxnvalue)))
"""Enumerated type to be used in tests."""

def append_columns(classdict, shape=()):
    """
    Append a ``Col`` of each PyTables data type to the `classdict`.

    A column of a certain TYPE gets called ``cTYPE``.  The number of
    added columns is returned.
    """
    for (itype, ptype) in enumerate(sorted(ptype_info.iterkeys())):
        colpos = itype + 1
        colname = 'c%s' % ptype
        colclass = getattr(tables, '%sCol' % ptype)
        colargs, colkwargs = [], {'shape': shape, 'pos': colpos}
        if ptype == 'Enum':
            colargs = [enum, enum(0), dtype_from_ptype[ptype]] + colargs
        elif ptype == 'String':
            colkwargs.update({'length': dtype_from_ptype[ptype].itemsize})
        classdict[colname] = colclass(*colargs, **colkwargs)
    ncols = colpos
    return ncols

def nested_description(classname, pos, shape=()):
    """
    Return a nested column description with all PyTables data types.

    A column of a certain TYPE gets called ``cTYPE``.  The nested
    column will be placed in the position indicated by `pos`.
    """
    classdict = {}
    append_columns(classdict, shape=shape)
    classdict['_v_pos'] = pos
    return new.classobj(classname, (tables.IsDescription,), classdict)

def table_description(classname, nclassname, shape=()):
    """
    Return a table description for testing queries.

    The description consists of all PyTables data types, both in the
    top level and in the ``cNested`` nested column.  A column of a
    certain TYPE gets called ``cTYPE``.  An extra integer column
    ``cExtra`` is also provided.  If a `shape` is given, it will be
    used for all columns.
    """
    classdict = {}
    colpos = append_columns(classdict, shape)

    ndescr = nested_description(nclassname, colpos, shape=shape)
    classdict['cNested'] = ndescr
    colpos += 1

    extracol = tables.IntCol(shape=shape, pos=colpos)
    classdict['cExtra'] = extracol
    colpos += 1

    return new.classobj(classname, (tables.IsDescription,), classdict)

TableDescription = table_description(
    'TableDescription', 'NestedDescription' )
"""Unidimensional table description for testing queries."""

MDTableDescription = table_description(
    'MDTableDescription', 'MDNestedDescription', shape=md_shape )
"""Multidimensional table description for testing queries."""

table_data = {}
"""Cached table data for a given shape and number of rows."""
# Data is cached because computing it row by row is quite slow.  Hop!

def fill_table(table, shape, nrows):
    """
    Fill the given `table` with `nrows` rows of data.

    Values in the i-th row (where 0 <= i < `row_period`) for a
    multidimensional field with M elements span from i to i+M-1.  For
    subsequent rows, values repeat cyclically.

    The same goes for the ``cExtra`` column, but values range from
    -`row_period`/2 to +`row_period`/2.
    """
    # Reuse already computed data if possible.
    tdata = table_data.get((shape, nrows))
    if tdata is not None:
        table.append(tdata)
        table.flush()
        return

    size = int(numpy.product(shape))

    row, value = table.row, 0
    for nrow in xrange(nrows):
        data = numpy.arange(value, value + size).reshape(shape)
        for (ptype, dtype) in dtype_from_ptype.iteritems():
            colname = 'c%s' % ptype
            ncolname = 'cNested/%s' % colname
            if ptype == 'Bool':
                coldata = data > (row_period / 2)
            elif ptype == 'String':
                sdata = [str_format % x for x in range(value, value + size)]
                coldata = numpy.array(sdata, dtype=dtype).reshape(shape)
            else:
                coldata = numpy.asarray(data, dtype=dtype)
            row[ncolname] = row[colname] = coldata
            row['cExtra'] = data - (row_period / 2)
        row.append()
        value += 1
        if value == row_period:
            value = 0
    table.flush()

    # Make computed data reusable.
    tdata = table.read()
    table_data[(shape, nrows)] = tdata

class TableQueryTestCase(tests.TempFileMixin, tests.PyTablesTestCase):

    """
    Test querying table data.

    Sub-classes must define the following attributes:

    ``tableDescription``
        The description of the table to be created.
    ``shape``
        The shape of data fields in the table.
    ``nrows``
        The number of data rows to be generated for the table.

    Sub-classes may redefine the following attributes:

    ``indexed``
        Whether columns shall be indexed, if possible.  Default is not
        to index them.
    ``optlevel``
        The level of optimisation of column indexes.  Default is 0.
    """

    indexed = False
    optlevel = 0

    colNotIndexable_re = re.compile(r"\bcan not be indexed\b")
    condNotBoolean_re = re.compile(r"\bdoes not have a boolean type\b")

    def createIndexes(self, colname, ncolname):
        if not self.indexed:
            return
        try:
            vprint("* Creating index on column ``%s``..." % colname, nonl=True)
            for acolname in [colname, ncolname]:
                acolumn = self.table.colinstances[acolname]
                acolumn.createIndex(optlevel=self.optlevel, testmode=True)
            vprint("ok.")
        except TypeError, te:
            if self.colNotIndexable_re.search(str(te)):
                vprint("can not be indexed.")
                raise tests.SkipTest  # can't be indexed, nothing new to test
            raise
        except NotImplementedError:
            vprint("not supported yet.")
            raise tests.SkipTest  # column does not support indexing yet

    def setUp(self):
        super(TableQueryTestCase, self).setUp()
        self.table = table = self.h5file.createTable(
            '/', 'test', self.tableDescription, expectedrows=self.nrows )
        fill_table(table, self.shape, self.nrows)

    ## XXX Need some standard checks on query usage.

operators = [None, '<', '==', '!=']
"""Comparison operators to check with different types."""
if tests.heavy:
    operators += ['~', '<=', '>=', '>']
## XXX Need to check por operator pairs.
left_bound = row_period / 4
"""Operand of left side operator in comparisons with operator pairs."""
right_bound = row_period * 3 / 4
"""Operand of right side operator in comparisons with operator pairs."""
extra_conditions = ['', '& ((cExtra+1) > 0)', '| ((cExtra+1) > 0)']
"""Extra conditions to append to comparison conditions."""

def create_test_method(ptype, op, extracond):
    bound = right_bound
    if ptype == 'String':
        bound = str_format % bound
    bound = nxtype_from_ptype[ptype](bound)
    condvars = {'bound': bound}

    colname = 'c%s' % ptype
    ncolname = 'cNested/%s' % colname
    if not op:
        cond = colname
    elif op == '~':
        cond = '~(%s)' % colname
    elif op == '<':
        cond = '%s %s %s' % (colname, op, bound)  # variable-constant
    else:
        cond = '%s %s bound' % (colname, op)  # variable-variable
    if extracond:
        cond = '(%s) %s' % (cond, extracond)

    def test_method(self):
        vprint("* Condition is ``%s``." % cond)
        # Replace bitwise operators with their logical counterparts.
        pycond = cond
        for (ptop, pyop) in [('&', 'and'), ('|', 'or'), ('~', 'not')]:
            pycond = pycond.replace(ptop, pyop)
        pycond = compile(pycond, '<string>', 'eval')

        table = self.table
        self.createIndexes(colname, ncolname)

        reflen = None
        ## XXX Better try to check retrieved data with ``readWhere()``.
        ## XXX Better yet, try all query methods in parallel!
        # Test that both simple and nested columns work as expected.
        # Because of the way the table is filled, results are the same.
        for acolname in [colname, ncolname]:
            # First the reference Python version.
            pylen, pyvars = 0, condvars.copy()
            for row in table:
                pyvars[colname] = row[acolname]
                pyvars['cExtra'] = row['cExtra']
                try:
                    isvalidrow = eval(pycond, {}, pyvars)
                except TypeError:
                    vprint("* Python type does not support the operation.")
                    raise tests.SkipTest
                if isvalidrow:
                    pylen += 1
            if reflen is None:  # initialise reference length
                reflen = pylen
            vprint( "* %d rows selected by Python from ``%s``."
                    % (pylen, acolname) )
            self.assertEqual(pylen, reflen)

            # Then the in-kernel or indexed version.
            ptvars = condvars.copy()
            ptvars[colname] = table.colinstances[acolname]
            ptvars['cExtra'] = table.colinstances['cExtra']
            try:
                ptlen = len([r for r in table.where(cond, ptvars)])
            except TypeError, te:
                if self.condNotBoolean_re.search(str(te)):
                    vprint("* Condition is not boolean.")
                    raise tests.SkipTest
                raise
            except NotImplementedError:
                vprint("* PyTables type does not support the operation.")
                raise tests.SkipTest
            vprint( "* %d rows selected by PyTables from ``%s``."
                    % (ptlen, acolname) )
            self.assertEqual(ptlen, reflen)

    test_method.__doc__ = "Testing ``%s``." % cond
    return test_method

# Create individual tests.  You may restrict which tests are generated
# by replacing the sequences in the ``for`` statements.  For instance:
testn = 0
for ptype in ptype_info:  # for ptype in ['String']:
    for op in operators:  # for op in ['!=']:
        for extracond in extra_conditions:  # for extracond in ['']:
            tmethod = create_test_method(ptype, op, extracond)
            tmethod.__name__ = 'test_a%04d' % testn
            ptmethod = tests.pyTablesTest(tmethod)
            imethod = new.instancemethod(ptmethod, None, TableQueryTestCase)
            setattr(TableQueryTestCase, tmethod.__name__, imethod)
            testn += 1

# Base classes for all queries.
class ScalarTableMixin:
    tableDescription = TableDescription
    shape = ()
class MDTableMixin:
    tableDescription = MDTableDescription
    shape = md_shape

# Base classes for non-indexed queries.
NX_BLOCK_SIZE1 = 128  # from ``interpreter.c`` in Numexpr
NX_BLOCK_SIZE2 = 8  # from ``interpreter.c`` in Numexpr

class SmallTableMixin:
    nrows = row_period * 2
    assert NX_BLOCK_SIZE2 < nrows < NX_BLOCK_SIZE1
    assert nrows % NX_BLOCK_SIZE2 != 0  # to have some residual rows
class BigTableMixin:
    nrows = row_period * 3
    assert nrows > NX_BLOCK_SIZE1 + NX_BLOCK_SIZE2
    assert nrows % NX_BLOCK_SIZE1 != 0
    assert nrows % NX_BLOCK_SIZE2 != 0  # to have some residual rows

# Parameters for non-indexed queries.
table_sizes = ['Small']
if tests.heavy:
    table_sizes += ['Big']
table_ndims = ['Scalar']  # to enable multidimensional testing, include 'MD'

# Non-indexed queries: ``[SB][SM]TQTestCase``.
def niclassdata():
    for size in table_sizes:
        for ndim in table_ndims:
            classname = '%s%sTQTestCase' % (size[0], ndim[0])
            cbasenames = ( '%sTableMixin' % size, '%sTableMixin' % ndim,
                           'TableQueryTestCase' )
            yield (classname, cbasenames, {})

# Base classes for indexed queries.
class SmallITableMixin:
    nrows = 50
class MediumITableMixin:
    nrows = 100
class BigITableMixin:
    nrows = 500

# Parameters for indexed queries.
itable_sizes = ['Small']
itable_optvalues = [0, 2, 3]
if tests.heavy:
    itable_sizes += ['Medium', 'Big']
    itable_optvalues += [7, 9]

# Indexed queries: ``[SMB]I[02379]TQTestCase``.
def iclassdata():
    for size in itable_sizes:
        for optlevel in itable_optvalues:
            classname = '%sI%dTQTestCase' % (size[0], optlevel)
            cbasenames = ( '%sITableMixin' % size, 'ScalarTableMixin',
                           'TableQueryTestCase' )
            yield ( classname, cbasenames,
                    {'optlevel': optlevel, 'indexed': True} )

# Create test classes.
for cdatafunc in [niclassdata, iclassdata]:
    for (cname, cbasenames, cdict) in cdatafunc():
        cbases = tuple(eval(cbase) for cbase in cbasenames)
        class_ = new.classobj(cname, cbases, cdict)
        exec '%s = class_' % cname

def suite():
    """Return a test suite consisting of all the test cases in the module."""

    testSuite = unittest.TestSuite()

    niter = 1
    for i in range(niter):
        for cdatafunc in [niclassdata, iclassdata]:
            for cdata in cdatafunc():
                cname = cdata[0]
                testSuite.addTest(unittest.makeSuite(eval(cname)))

    return testSuite

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


<commit_before>"""
Functions dealing with Modified Preorder Tree Traversal related setup
and creation of instance methods for a model, given the names of its
tree attributes.

``treeify`` should be the only function a user of this application need
use directly to set their model up for Modified Preorder Tree Traversal.
"""
from django.db import models
from django.db.models import signals
from django.dispatch import dispatcher

from mptt.signals import pre_delete, pre_save
from mptt.managers import TreeManager

__all__ = ['treeify']

def treeify(cls, parent_attr='parent', left_attr='lft', right_attr='rght',
            tree_id_attr='tree_id', level_attr='level',
            tree_manager_attr='tree'):
    """
    Sets the given model class up for Modified Preorder Tree Traversal,
    which involves:

    1. If any of the specified tree fields -- ``left_attr``,
       ``right_attr``, ``tree_id_attr`` and ``level_attr`` -- do not
       exist, adding them to the model class dynamically.
    2. Creating pre_save and pre_delete signal receiving functions to
       manage tree field contents.
    3. Adding tree related instance methods to the model class.
    4. Adding a custom tree ``Manager`` to the model class.
    """
    # Add tree fields if they do not exist
    for attr in [left_attr, right_attr, tree_id_attr, level_attr]:
        try:
            cls._meta.get_field(attr)
        except models.FieldDoesNotExist:
            models.PositiveIntegerField(
                db_index=True, editable=False).contribute_to_class(cls, attr)
    # Specifying weak=False is required in this case as the dispatcher
    # will be the only place a reference is held to the signal receiving
    # functions we're creating.
    dispatcher.connect(
        pre_save(parent_attr, left_attr, right_attr, tree_id_attr, level_attr),
        signal=signals.pre_save, sender=cls, weak=False)
    dispatcher.connect(pre_delete(left_attr, right_attr, tree_id_attr),
                       signal=signals.pre_delete, sender=cls, weak=False)
    setattr(cls, 'get_ancestors',
            get_ancestors(parent_attr, left_attr, right_attr, tree_id_attr))
    setattr(cls, 'get_descendants',
            get_descendants(left_attr, right_attr, tree_id_attr))
    setattr(cls, 'get_descendant_count',
            get_descendant_count(left_attr, right_attr))
    setattr(cls, 'get_next_sibling',
            get_next_sibling(parent_attr, left_attr, right_attr, tree_id_attr))
    setattr(cls, 'get_previous_sibling',
            get_previous_sibling(parent_attr, left_attr, right_attr,
                                 tree_id_attr))
    setattr(cls, 'get_siblings',
            get_siblings(parent_attr, left_attr, tree_id_attr))
    setattr(cls, 'is_child_node', is_child_node)
    setattr(cls, 'is_root_node', is_root_node(parent_attr))
    setattr(cls, 'move_to', move_to)
    if not hasattr(cls, tree_manager_attr):
        TreeManager(parent_attr, left_attr, right_attr, tree_id_attr,
                    level_attr).contribute_to_class(cls, tree_manager_attr)
    setattr(cls, '_tree_manager', getattr(cls, tree_manager_attr))

def get_ancestors(parent_attr, left_attr, right_attr, tree_id_attr):
    """
    Creates a function which retrieves the ancestors of a model instance
    which has the given tree attributes.
    """
    def _get_ancestors(instance, ascending=False):
        """
        Creates a ``QuerySet`` containing all the ancestors of this
        model instance.

        This defaults to being in descending order (root ancestor first,
        immediate parent last); passing ``True`` for the ``ascending``
        argument will reverse the ordering (immediate parent first, root
        ancestor last).
        """
        if getattr(instance, parent_attr) is None:
            return instance._default_manager.none()
        else:
            return instance._default_manager.filter(**{
                '%s__lt' % left_attr: getattr(instance, left_attr),
                '%s__gt' % right_attr: getattr(instance, right_attr),
                tree_id_attr: getattr(instance, tree_id_attr),
            }).order_by('%s%s' % ({True: '-', False: ''}[ascending], left_attr))
    return _get_ancestors

def get_descendants(left_attr, right_attr, tree_id_attr):
    """
    Creates a function which retrieves descendants of a model
    instance which has the given tree attributes.
    """
    def _get_descendants(instance, include_self=False):
        """
        Creates a ``QuerySet`` containing descendants of this model
        instance.

        If ``include_self`` is ``True``, the ``QuerySet`` will also
        include this model instance.
        """
        filters = {tree_id_attr: getattr(instance, tree_id_attr)}
        if include_self:
            filters['%s__range' % left_attr] = (getattr(instance, left_attr),
                                                getattr(instance, right_attr))
        else:
            filters['%s__gt' % left_attr] = getattr(instance, left_attr)
            filters['%s__lt' % left_attr] = getattr(instance, right_attr)
        return instance._default_manager.filter(**filters).order_by(left_attr)
    return _get_descendants

def get_descendant_count(left_attr, right_attr):
    """
    Creates a function which determines the number of descendants of a
    model instance which has the given tree attributes.
    """
    def _get_descendant_count(instance):
        """
        Returns the number of descendants this model instance has.
        """
        return (getattr(instance, right_attr) - getattr(instance, left_attr) - 1) / 2
    return _get_descendant_count

def get_previous_sibling(parent_attr, left_attr, right_attr, tree_id_attr):
    """
    Creates a function which retrieves the previous sibling of a model
    instance which has the given tree attributes.
    """
    def _get_previous_sibling(instance):
        """
        Returns this model instance's previous sibling in the tree, or
        ``None`` if it doesn't have a previous sibling.
        """
        if instance.is_root_node():
            filters = {
                '%s__isnull' % parent_attr: True,
                '%s__lt' % tree_id_attr: getattr(instance, tree_id_attr),
            }
            order_by = '-%s' % tree_id_attr
        else:
            filters = {
                 parent_attr: getattr(instance, '%s_id' % parent_attr),
                '%s__lt' % right_attr: getattr(instance, left_attr),
            }
            order_by = '-%s' % right_attr

        sibling = None
        try:
            sibling = instance._default_manager.filter(**filters).order_by(order_by)[0]
        except IndexError:
            pass
        return sibling
    return _get_previous_sibling

def get_next_sibling(parent_attr, left_attr, right_attr, tree_id_attr):
    """
    Creates a function which retrieves the next sibling of a model
    instance which has the given tree attributes.
    """
    def _get_next_sibling(instance):
        """
        Returns this model instance's next sibling in the tree, or
        ``None`` if it doesn't have a next sibling.
        """
        if instance.is_root_node():
            filters = {
                '%s__isnull' % parent_attr: True,
                '%s__gt' % tree_id_attr: getattr(instance, tree_id_attr),
            }
            order_by = tree_id_attr
        else:
            filters = {
                 parent_attr: getattr(instance, '%s_id' % parent_attr),
                '%s__gt' % left_attr: getattr(instance, right_attr),
            }
            order_by = left_attr

        sibling = None
        try:
            sibling = instance._default_manager.filter(**filters).order_by(order_by)[0]
        except IndexError:
            pass
        return sibling
    return _get_next_sibling

def get_siblings(parent_attr, left_attr, tree_id_attr):
    """
    Creates a function which retrieves siblings of a model instance
    which has the given tree attributes.
    """
    def _get_siblings(instance, include_self=False):
        """
        Creates a ``QuerySet`` containing siblings of this model
        instance. Root nodes are considered to be siblings of other root
        nodes.

        If ``include_self`` is ``True``, the ``QuerySet`` will also
        include this model instance.
        """
        if instance.is_root_node():
            queryset = instance._default_manager.filter(**{
                '%s__isnull' % parent_attr: True,
            }).order_by(tree_id_attr)
        else:
            queryset = instance._default_manager.filter(**{
                parent_attr: getattr(instance, '%s_id' % parent_attr),
            }).order_by(left_attr)
        if not include_self:
            queryset = queryset.exclude(pk=instance.pk)
        return queryset
    return _get_siblings

def is_child_node(instance):
    """
    Returns ``True`` if this model instance is a child node, ``False``
    otherwise.
    """
    return not instance.is_root_node()

def is_root_node(parent_attr):
    """
    Creates a function which determines if a model instance which has
    the given tree attributes is a root node.
    """
    def _is_root_node(instance):
        """
        Returns ``True`` if this model instance is a root node,
        ``False`` otherwise.
        """
        return getattr(instance, '%s_id' % parent_attr) is None
    return _is_root_node

def move_to(instance, target, position='first-child'):
    """
    Convenience method for calling ``TreeManager.move_to`` with this
    model instance.
    """
    instance._tree_manager.move_node(instance, target, position)
<commit_msg>Saved a query in get_ancestors by using is_root_node instead of retrieving the parent; saved a query for a specific case in get_descendants by returning an empty queryset<commit_after>"""
Functions dealing with Modified Preorder Tree Traversal related setup
and creation of instance methods for a model, given the names of its
tree attributes.

``treeify`` should be the only function a user of this application need
use directly to set their model up for Modified Preorder Tree Traversal.
"""
from django.db import models
from django.db.models import signals
from django.dispatch import dispatcher

from mptt.signals import pre_delete, pre_save
from mptt.managers import TreeManager

__all__ = ['treeify']

def treeify(cls, parent_attr='parent', left_attr='lft', right_attr='rght',
            tree_id_attr='tree_id', level_attr='level',
            tree_manager_attr='tree'):
    """
    Sets the given model class up for Modified Preorder Tree Traversal,
    which involves:

    1. If any of the specified tree fields -- ``left_attr``,
       ``right_attr``, ``tree_id_attr`` and ``level_attr`` -- do not
       exist, adding them to the model class dynamically.
    2. Creating pre_save and pre_delete signal receiving functions to
       manage tree field contents.
    3. Adding tree related instance methods to the model class.
    4. Adding a custom tree ``Manager`` to the model class.
    """
    # Add tree fields if they do not exist
    for attr in [left_attr, right_attr, tree_id_attr, level_attr]:
        try:
            cls._meta.get_field(attr)
        except models.FieldDoesNotExist:
            models.PositiveIntegerField(
                db_index=True, editable=False).contribute_to_class(cls, attr)
    # Specifying weak=False is required in this case as the dispatcher
    # will be the only place a reference is held to the signal receiving
    # functions we're creating.
    dispatcher.connect(
        pre_save(parent_attr, left_attr, right_attr, tree_id_attr, level_attr),
        signal=signals.pre_save, sender=cls, weak=False)
    dispatcher.connect(pre_delete(left_attr, right_attr, tree_id_attr),
                       signal=signals.pre_delete, sender=cls, weak=False)
    setattr(cls, 'get_ancestors',
            get_ancestors(parent_attr, left_attr, right_attr, tree_id_attr))
    setattr(cls, 'get_descendants',
            get_descendants(left_attr, right_attr, tree_id_attr))
    setattr(cls, 'get_descendant_count',
            get_descendant_count(left_attr, right_attr))
    setattr(cls, 'get_next_sibling',
            get_next_sibling(parent_attr, left_attr, right_attr, tree_id_attr))
    setattr(cls, 'get_previous_sibling',
            get_previous_sibling(parent_attr, left_attr, right_attr,
                                 tree_id_attr))
    setattr(cls, 'get_siblings',
            get_siblings(parent_attr, left_attr, tree_id_attr))
    setattr(cls, 'is_child_node', is_child_node)
    setattr(cls, 'is_root_node', is_root_node(parent_attr))
    setattr(cls, 'move_to', move_to)
    if not hasattr(cls, tree_manager_attr):
        TreeManager(parent_attr, left_attr, right_attr, tree_id_attr,
                    level_attr).contribute_to_class(cls, tree_manager_attr)
    setattr(cls, '_tree_manager', getattr(cls, tree_manager_attr))

def get_ancestors(parent_attr, left_attr, right_attr, tree_id_attr):
    """
    Creates a function which retrieves the ancestors of a model instance
    which has the given tree attributes.
    """
    def _get_ancestors(instance, ascending=False):
        """
        Creates a ``QuerySet`` containing all the ancestors of this
        model instance.

        This defaults to being in descending order (root ancestor first,
        immediate parent last); passing ``True`` for the ``ascending``
        argument will reverse the ordering (immediate parent first, root
        ancestor last).
        """
        if instance.is_root_node():
            return instance._default_manager.none()
        else:
            return instance._default_manager.filter(**{
                '%s__lt' % left_attr: getattr(instance, left_attr),
                '%s__gt' % right_attr: getattr(instance, right_attr),
                tree_id_attr: getattr(instance, tree_id_attr),
            }).order_by('%s%s' % ({True: '-', False: ''}[ascending], left_attr))
    return _get_ancestors

def get_descendants(left_attr, right_attr, tree_id_attr):
    """
    Creates a function which retrieves descendants of a model
    instance which has the given tree attributes.
    """
    def _get_descendants(instance, include_self=False):
        """
        Creates a ``QuerySet`` containing descendants of this model
        instance.

        If ``include_self`` is ``True``, the ``QuerySet`` will also
        include this model instance.
        """
        if not include_self and not instance.get_descendant_count():
            return instance._default_manager.none()

        filters = {tree_id_attr: getattr(instance, tree_id_attr)}
        if include_self:
            filters['%s__range' % left_attr] = (getattr(instance, left_attr),
                                                getattr(instance, right_attr))
        else:
            filters['%s__gt' % left_attr] = getattr(instance, left_attr)
            filters['%s__lt' % left_attr] = getattr(instance, right_attr)
        return instance._default_manager.filter(**filters).order_by(left_attr)
    return _get_descendants

def get_descendant_count(left_attr, right_attr):
    """
    Creates a function which determines the number of descendants of a
    model instance which has the given tree attributes.
    """
    def _get_descendant_count(instance):
        """
        Returns the number of descendants this model instance has.
        """
        return (getattr(instance, right_attr) - getattr(instance, left_attr) - 1) / 2
    return _get_descendant_count

def get_previous_sibling(parent_attr, left_attr, right_attr, tree_id_attr):
    """
    Creates a function which retrieves the previous sibling of a model
    instance which has the given tree attributes.
    """
    def _get_previous_sibling(instance):
        """
        Returns this model instance's previous sibling in the tree, or
        ``None`` if it doesn't have a previous sibling.
        """
        if instance.is_root_node():
            filters = {
                '%s__isnull' % parent_attr: True,
                '%s__lt' % tree_id_attr: getattr(instance, tree_id_attr),
            }
            order_by = '-%s' % tree_id_attr
        else:
            filters = {
                 parent_attr: getattr(instance, '%s_id' % parent_attr),
                '%s__lt' % right_attr: getattr(instance, left_attr),
            }
            order_by = '-%s' % right_attr

        sibling = None
        try:
            sibling = instance._default_manager.filter(**filters).order_by(order_by)[0]
        except IndexError:
            pass
        return sibling
    return _get_previous_sibling

def get_next_sibling(parent_attr, left_attr, right_attr, tree_id_attr):
    """
    Creates a function which retrieves the next sibling of a model
    instance which has the given tree attributes.
    """
    def _get_next_sibling(instance):
        """
        Returns this model instance's next sibling in the tree, or
        ``None`` if it doesn't have a next sibling.
        """
        if instance.is_root_node():
            filters = {
                '%s__isnull' % parent_attr: True,
                '%s__gt' % tree_id_attr: getattr(instance, tree_id_attr),
            }
            order_by = tree_id_attr
        else:
            filters = {
                 parent_attr: getattr(instance, '%s_id' % parent_attr),
                '%s__gt' % left_attr: getattr(instance, right_attr),
            }
            order_by = left_attr

        sibling = None
        try:
            sibling = instance._default_manager.filter(**filters).order_by(order_by)[0]
        except IndexError:
            pass
        return sibling
    return _get_next_sibling

def get_siblings(parent_attr, left_attr, tree_id_attr):
    """
    Creates a function which retrieves siblings of a model instance
    which has the given tree attributes.
    """
    def _get_siblings(instance, include_self=False):
        """
        Creates a ``QuerySet`` containing siblings of this model
        instance. Root nodes are considered to be siblings of other root
        nodes.

        If ``include_self`` is ``True``, the ``QuerySet`` will also
        include this model instance.
        """
        if instance.is_root_node():
            queryset = instance._default_manager.filter(**{
                '%s__isnull' % parent_attr: True,
            }).order_by(tree_id_attr)
        else:
            queryset = instance._default_manager.filter(**{
                parent_attr: getattr(instance, '%s_id' % parent_attr),
            }).order_by(left_attr)
        if not include_self:
            queryset = queryset.exclude(pk=instance.pk)
        return queryset
    return _get_siblings

def is_child_node(instance):
    """
    Returns ``True`` if this model instance is a child node, ``False``
    otherwise.
    """
    return not instance.is_root_node()

def is_root_node(parent_attr):
    """
    Creates a function which determines if a model instance which has
    the given tree attributes is a root node.
    """
    def _is_root_node(instance):
        """
        Returns ``True`` if this model instance is a root node,
        ``False`` otherwise.
        """
        return getattr(instance, '%s_id' % parent_attr) is None
    return _is_root_node

def move_to(instance, target, position='first-child'):
    """
    Convenience method for calling ``TreeManager.move_to`` with this
    model instance.
    """
    instance._tree_manager.move_node(instance, target, position)
<|endoftext|>