def test_init_from_object(self):
query1 = Query.from_object([1, 3, 4, 2])
self.assertEqual(query1.source, [1, 3, 4, 2])
self.assertEqual(query1.args, ())
self.assertEqual(query1.kwds, {})
self.assertEqual(query1._query_steps, [])
query2 = Query.from_object({'a': 1, 'b': 2})
self.assertEqual(query2.source, {'a': 1, 'b': 2})
self.assertEqual(query2.args, ())
self.assertEqual(query2.kwds, {})
self.assertEqual(query2._query_steps, [])
# When from_object() receives a Query, it should return
# a copy rather than trying to use it as a data object.
query3 = Query.from_object(query2)
self.assertIsNot(query3, query2)
self.assertEqual(query3.source, {'a': 1, 'b': 2})
self.assertEqual(query3.args, ())
self.assertEqual(query3.kwds, {})
self.assertEqual(query3._query_steps, [])
query4 = Query.from_object('abc')
self.assertEqual(query4.source, ['abc'], msg=\
'Strings or non-iterables should be wrapped as a list')
self.assertEqual(query4.args, ())
self.assertEqual(query4.kwds, {})
self.assertEqual(query4._query_steps, [])
query5 = Query.from_object(123)
self.assertEqual(query5.source, [123], msg=\
'Strings or non-iterables should be wrapped as a list')
self.assertEqual(query5.args, ())
self.assertEqual(query5.kwds, {})
self.assertEqual(query5._query_steps, [])
def test_execute_other_source(self):
query = Query.from_object([1, 3, 4, 2])
result = query.execute()
self.assertIsInstance(result, Result)
self.assertEqual(result.fetch(), [1, 3, 4, 2])
query = Query.from_object(iter([1, 3, 4, 2]))
result = query.execute()
self.assertIsInstance(result, Result)
self.assertEqual(result.fetch(), [1, 3, 4, 2])
query = Query.from_object(Result([1, 3, 4, 2], evaltype=list))
result = query.execute()
self.assertIsInstance(result, Result)
self.assertEqual(result.fetch(), [1, 3, 4, 2])
query = Query.from_object(Result({'a': 1, 'b': 2}, evaltype=dict))
result = query.execute()
self.assertIsInstance(result, Result)
self.assertEqual(result.fetch(), {'a': 1, 'b': 2})
query = Query.from_object(IterItems(iter([iter(['a', 1]), iter(['b', 2])])))
result = query.execute()
self.assertIsInstance(result, Result)
self.assertEqual(result.fetch(), {'a': 1, 'b': 2})
def test_flatten(self):
query1 = Query({'col1': ('col2', 'col2')})
query2 = query1.flatten()
self.assertIsNot(query1, query2, 'should return new object')
source = Select([('col1', 'col2'), ('a', '2'), ('b', '2')])
result = query2.execute(source)
self.assertEqual(result.fetch(), [('a', '2', '2'), ('b', '2', '2')])
def test_starmap(self):
query1 = Query([('col2', 'col2')])
query2 = query1.starmap(lambda x, y: x + y)
self.assertIsNot(query1, query2, 'should return new object')
source = Select([('col1', 'col2'), ('a', 1), ('b', 2)])
result = query2.execute(source)
self.assertEqual(result.fetch(), [2, 4])
def test_map(self):
query1 = Query(['col2'])
query2 = query1.map(int)
self.assertIsNot(query1, query2, 'should return new object')
source = Select([('col1', 'col2'), ('a', '2'), ('b', '2')])
result = query2.execute(source)
self.assertEqual(result.fetch(), [2, 2])
def test_unwrap(self):
query1 = Query({'col1': ['col2']})
query2 = query1.unwrap()
self.assertIsNot(query1, query2, 'should return new object')
source = Select([('col1', 'col2'), ('a', 1), ('b', 2), ('b', 3)])
result = query2.execute(source)
self.assertEqual(result.fetch(), {'a': 1, 'b': [2, 3]})
def test_init_with_select(self):
source = Select([('A', 'B'), (1, 2), (1, 2)])
query = Query(source, ['A'], B=2)
self.assertEqual(query.source, source)
self.assertEqual(query.args, (['A'],))
self.assertEqual(query.kwds, {'B': 2})
self.assertEqual(query._query_steps, [])
with self.assertRaises(TypeError):
query = Query(None, ['foo'], bar='baz')
def test_fetch_datasource(self):
select = Select([('A', 'B'), ('1', '2'), ('1', '2')])
query = Query(select, ['B'])
query._query_steps = [
('map', (int,), {}),
('map', (lambda x: x * 2,), {}),
('sum', (), {}),
]
result = query.fetch()
self.assertEqual(result, 8)
def test_explain(self):
query = Query(['col1'])
expected = """
Data Source:
<none given> (assuming Select object)
Execution Plan:
getattr, (<RESULT>, '_select'), {}
<RESULT>, (['col1']), {}
"""
expected = textwrap.dedent(expected).strip()
self.assertEqual(query._explain(file=None), expected)
def test_init_with_invalid_args(self):
# Missing args.
with self.assertRaises(TypeError, msg='should require select args'):
Query()
# Bad "select" field.
source = Select([('A', 'B'), (1, 2), (1, 2)])
with self.assertRaises(LookupError, msg='should fail immediately when fieldname conflicts with provided source'):
query = Query(source, ['X'], B=2)
# Bad "where" field.
source = Select([('A', 'B'), (1, 2), (1, 2)])
with self.assertRaises(LookupError, msg='should fail immediately when fieldname conflicts with provided "where" field'):
query = Query(source, ['A'], Y=2)
def test_bad_truncation(self):
"""Should not get truncated by preview_length."""
select = Select([
('A', 'B'),
('x', 1),
('x', 2),
('x', 3),
('x', 4),
('x', 5),
('x', 6),
('x', 7),
('x', 8),
('x', 9),
])
query = Query(select, {'A': 'B'})
self.assertEqual(query.fetch(), {'x': [1, 2, 3, 4, 5, 6, 7, 8, 9]})
def test_init_with_nested_dicts(self):
"""Support for nested dictionaries was removed (for now).
It's likely that arbitrary nesting would complicate the ability
to check complex data values that are, themselves, mappings
(like probability mass functions represented as a dictionary).
"""
regex = 'mappings can not be nested'
with self.assertRaisesRegex(ValueError, regex):
query = Query({'A': {'B': 'C'}}, D='x')
def test_execute_datasource(self):
select = Select([('A', 'B'), ('1', '2'), ('1', '2')])
query = Query(select, ['B'])
query._query_steps = [
('map', (int,), {}),
('map', (lambda x: x * 2,), {}),
('sum', (), {}),
]
result = query.execute()
self.assertEqual(result, 8)
query = Query(['A'])
regex = "expected 'Select', got 'list'"
with self.assertRaisesRegex(TypeError, regex):
query.execute(['hello', 'world']) # <- Expects None or Query, not list!
def test_explain2(self):
query = Query(['label1'])
expected = """
Data Source:
<none given> (assuming Select object)
Execution Plan:
getattr, (<RESULT>, '_select'), {}
<RESULT>, (['label1']), {}
"""
expected = textwrap.dedent(expected).strip()
# Defaults to stdout (redirected to StringIO for testing).
string_io = StringIO()
returned_value = query._explain(file=string_io)
self.assertIsNone(returned_value)
printed_value = string_io.getvalue().strip()
self.assertEqual(printed_value, expected)
# Get result as string.
returned_value = query._explain(file=None)
self.assertEqual(returned_value, expected)
def test_repr(self):
# Check "no select" signature.
query = Query(['label1'])
regex = r"Query\(\[u?'label1'\]\)"
self.assertRegex(repr(query), regex)
# Check "no select" with keyword string.
query = Query(['label1'], label2='x')
regex = r"Query\(\[u?'label1'\], label2='x'\)"
self.assertRegex(repr(query), regex)
# Check "no select" with keyword list.
query = Query(['label1'], label2=['x', 'y'])
regex = r"Query\(\[u?'label1'\], label2=\[u?'x', u?'y'\]\)"
self.assertRegex(repr(query), regex)
# Check "select-provided" signature.
select = Select([('A', 'B'), ('x', 1), ('y', 2), ('z', 3)])
query = Query(select, ['B'])
short_repr = super(Select, select).__repr__()
expected = "Query({0}, {1!r})".format(short_repr, ['B'])
#print(repr(query))
self.assertEqual(repr(query), expected)
# Check "from_object" signature.
query = Query.from_object([1, 2, 3])
expected = "Query.from_object([1, 2, 3])"
self.assertEqual(repr(query), expected)
# Check query steps.
query = Query(['label1']).distinct().count()
regex = r"Query\(\[u?'label1'\]\).distinct\(\).count\(\)"
self.assertRegex(repr(query), regex)
# Check query steps with function argument.
def upper(x):
return str(x.upper())
query = Query(['label1']).map(upper)
regex = r"Query\(\[u?'label1'\]\).map\(upper\)"
self.assertRegex(repr(query), regex)
# Check query steps with lambda argument.
lower = lambda x: str(x).lower()
query = Query(['label1']).map(lower)
regex = r"Query\(\[u?'label1'\]\).map\(<lambda>\)"
self.assertRegex(repr(query), regex)
def test_reduce(self):
query1 = Query.from_object({'a': [1, 3, 5], 'b': [2, 4, 6]})
# Test simple case.
query2 = query1.reduce(lambda x, y: x + y)
self.assertEqual(query2.fetch(), {'a': 9, 'b': 12})
# Test optional initializer_factory.
def func(acc, upd):
acc.append(upd)
return acc
query3 = query1.reduce(func, initializer_factory=list)
self.assertEqual(query3.fetch(), {'a': [1, 3, 5], 'b': [2, 4, 6]})
# Test bad initializer_factory.
with self.assertRaises(TypeError):
query4 = query1.reduce(func, initializer_factory=[])
def test_filter(self):
query1 = Query(['col1'])
query2 = query1.filter(lambda x: x == 'a')
self.assertIsNot(query1, query2, 'should return new object')
source = Select([('col1', 'col2'), ('a', '2'), ('b', '2')])
result = query2.execute(source)
self.assertEqual(result.fetch(), ['a'])
# No filter arg should default to bool()
source = Select([('col1',), (1,), (2,), (0,), (3,)])
query = Query(set(['col1'])).filter() # <- No arg!
result = query.execute(source)
self.assertEqual(result.fetch(), set([1, 2, 3]))
def test_optimize_distinct(self):
"""
Unoptimized:
Select._select({'col1': ['values']}, col2='xyz').distinct()
Optimized:
Select._select_distinct({'col1': ['values']}, col2='xyz')
"""
unoptimized = (
(getattr, (RESULT_TOKEN, '_select'), {}),
(RESULT_TOKEN, ({'col1': ['values']},), {'col2': 'xyz'}),
(_sqlite_distinct, (RESULT_TOKEN,), {}),
)
optimized = Query._optimize(unoptimized)
expected = (
(getattr, (RESULT_TOKEN, '_select_distinct'), {}),
(RESULT_TOKEN, ({'col1': ['values']},), {'col2': 'xyz'}),
)
self.assertEqual(optimized, expected)
def test_optimize_aggregation(self):
"""
Unoptimized:
Select._select({'col1': ['values']}, col2='xyz').sum()
Optimized:
Select._select_aggregate('SUM', {'col1': ['values']}, col2='xyz')
"""
unoptimized = (
(getattr, (RESULT_TOKEN, '_select'), {}),
(RESULT_TOKEN, ({'col1': ['values']},), {'col2': 'xyz'}),
(_apply_to_data, (_sqlite_sum, RESULT_TOKEN,), {}),
)
optimized = Query._optimize(unoptimized)
expected = (
(getattr, (RESULT_TOKEN, '_select_aggregate'), {}),
(RESULT_TOKEN, ('SUM', {'col1': ['values']},), {'col2': 'xyz'}),
)
self.assertEqual(optimized, expected)
def test__copy__(self):
# Select-arg only.
query = Query(['B'])
copied = query.__copy__()
self.assertIs(copied.source, query.source)
self.assertEqual(copied.args, query.args)
self.assertEqual(copied.kwds, query.kwds)
self.assertEqual(copied._query_steps, query._query_steps)
self.assertIsNot(copied.kwds, query.kwds)
self.assertIsNot(copied._query_steps, query._query_steps)
# Select and keyword.
query = Query(['B'], C='x')
copied = query.__copy__()
self.assertIs(copied.source, query.source)
self.assertEqual(copied.args, query.args)
self.assertEqual(copied.kwds, query.kwds)
self.assertEqual(copied._query_steps, query._query_steps)
# Source, columns, and keyword.
source = Select([('A', 'B'), (1, 2), (1, 2)])
query = Query(source, ['B'])
copied = query.__copy__()
self.assertIs(copied.source, query.source)
self.assertEqual(copied.args, query.args)
self.assertEqual(copied.kwds, query.kwds)
self.assertEqual(copied._query_steps, query._query_steps)
# Select and additional query methods.
query = Query(['B']).map(lambda x: str(x).upper())
copied = query.__copy__()
self.assertIs(copied.source, query.source)
self.assertEqual(copied.args, query.args)
self.assertEqual(copied.kwds, query.kwds)
self.assertEqual(copied._query_steps, query._query_steps)
def test_init_no_data(self):
# Use column and where syntax.
query = Query(['foo'], bar='baz')
self.assertEqual(query.source, None)
# Test query steps.
query = Query(['foo'], bar='baz')
self.assertEqual(query._query_steps, [])
# Adding query steps.
query = query.distinct().sum()
expected = [
('distinct', (), {}),
('sum', (), {}),
]
self.assertEqual(query._query_steps, expected)
# Single-string defaults to list-of-single-string.
query = Query('foo')
self.assertEqual(query.args[0], ['foo'], 'should be wrapped as list')
# Multi-item-container defaults to list-of-container.
query = Query(['foo', 'bar'])
self.assertEqual(query.args[0], [['foo', 'bar']], 'should be wrapped as list')
# Mapping with single-string defaults to list-of-single-string.
query = Query({'foo': 'bar'})
self.assertEqual(query.args[0], {'foo': ['bar']}, 'value should be wrapped as list')
# Mapping with multi-item-container defaults to list-of-container.
query = Query({'foo': ['bar', 'baz']})
self.assertEqual(query.args[0], {'foo': [['bar', 'baz']]}, 'value should be wrapped as list')
def test_iterate_single_result(self):
"""Single items should be wrapped as iterators when iterated over."""
select = Select([('A', 'B'), (1, 2), (1, 2)])
query = Query(select, ['B']).sum()
self.assertEqual(list(query), [4])
def test_iterate_source(self):
select = Select([('A', 'B'), (1, 2), (1, 2)])
query = Query(select, ['B'])
self.assertEqual(list(query), [2, 2])