Mentions légales du service

Skip to content
Snippets Groups Projects
Commit 58fc3d6c authored by Bruno Guillon's avatar Bruno Guillon
Browse files

Select: project/set_columns/set_column_origins/add_column

parent 8dc733c6
No related branches found
Tags dead
1 merge request!46Draft: Resolve "Project improved and set_columns"
Pipeline #991345 failed
......@@ -234,11 +234,11 @@ class KeyedTuple(tuple[E, ...], Generic[E]):
...
@overload
def __getitem__(self, i: list[int], /) -> list[E]:
def __getitem__(self, i: list[int | str], /) -> list[E]:
...
def __getitem__(
self, i: SupportsIndex | slice | list[int], /
self, i: SupportsIndex | slice | list[int | str], /
) -> E | KeyedTuple[E] | list[E]:
if isinstance(i, list):
# return list rather than self._store.__class__ to allow duplicates
......
from __future__ import annotations
import functools
import collections
from typing import Union, Iterable, Tuple, Optional, Mapping, Sequence
from typing import Any, Union, Iterable, Tuple, Optional, Mapping, Sequence
import querybuilder.classes.stores as qbstores
import querybuilder.utils.constants as qbconstants
from querybuilder.utils.decorators import method_accepting_lambdas
......@@ -212,7 +212,7 @@ class Select(DQLQuery):
return self.set_aliases(d)
def reset_aliases(self):
return self.buildfrom(self, aliases=())
return self.set_aliases(())
@method_accepting_lambdas
def set_where(self, where):
......@@ -248,20 +248,70 @@ class Select(DQLQuery):
def set_offset(self, offset):
return self.buildfrom(self, offset=offset)
def project(self, indices: Union[slice, Iterable[int]]) -> Select:
columns_and_aliases = [
(c, self.aliases.get(i, None)) for i, c in enumerate(self.column_origins)
]
def project(self, columns: Iterable[int | str] | slice) -> Select:
"""Project, possibly permuting, the columns of the relation
if isinstance(indices, slice):
columns_and_aliases = columns_and_aliases[indices]
else:
columns_and_aliases = [columns_and_aliases[i] for i in indices]
Parameters
----------
# TODO
columns = [e[0] for e in columns_and_aliases]
aliases = {i: e[1] for i, e in enumerate(columns_and_aliases) if e[1]}
Notes
-----
Projected columns should be columns of the relation, and cannot be repeated.
Aliases are preserved.
"""
if isinstance(columns, slice):
return self.project(range(len(self.columns))[columns])
return self.set_columns(columns, disallow_repetitions=True)
def set_columns(
self,
columns: Iterable[qbcolumns.Column | int | str],
*,
disallow_repetitions: bool = False,
) -> Select:
""" """
if self.distinct == True:
# TODO: define a DISTINCT ON clause on previously selected columns
raise NotImplementedError()
indices = set()
cols = []
aliases = {}
for i, c in enumerate(columns):
if isinstance(c, int):
j = c
c = self.column_origins[j]
else:
if isinstance(c, str):
c = self.column_origins._[c]
j = self.column_origins.index(c)
if disallow_repetitions and j in indices:
raise ValueError(f"Got repeated column {self.column_origins[j]}")
if j in self.aliases:
aliases[i] = self.aliases[j]
cols.append(c)
indices.add(j)
return self.set_column_origins(cols).set_aliases(aliases)
def set_column_origins(self, columns: Iterable[qbcolumns.Column | Any]) -> Select:
cols = [
c if isinstance(c, qbcolumns.Column) else qbcolumns.make_column(c)
for c in columns
]
return self.reset_aliases().buildfrom(self, columns=cols)
return self._set_columns(columns).set_aliases(aliases)
def add_column(
self, col: qbcolumns.Column | Any, /, *, alias: Optional[str] = None
) -> Select:
if not isinstance(col, qbcolumns.Column):
col = qbcolumns.make_column(col)
aliases = self.aliases
if alias:
aliases |= {len(self.columns): alias}
return self.set_column_origins(self.columns + (col,)).set_aliases(aliases)
@method_accepting_lambdas
def add_where(self, where):
......@@ -339,9 +389,6 @@ class Select(DQLQuery):
name, self, schema_name=schema_name, aliases=aliases, **kwargs
)
def _set_columns(self, columns):
return self.buildfrom(self, columns=columns)
def _post_getstate(self, state):
state["columns"] = state.pop("column_origins")
return state
......
......@@ -7,20 +7,46 @@ import querybuilder.queries.dql as dql
class TestSelect:
tk = qb.drivers.sql.tokenizer.Tokenizer()
def test_project_only_columns_with_list(self):
@pytest.mark.parametrize(
"projlist, explist",
[
([0, 4, 2], ["c0", "c4", "c2"]),
([0], ["c0"]),
([], []),
(["c0", "c4", "c2"], ["c0", "c4", "c2"]),
(["c0"], ["c0"]),
([0, "c4", 2], ["c0", "c4", "c2"]),
(slice(0, 4, 2), ["c0", "c2"]),
],
)
def test_project_only_columns_with_list(self, projlist, explist):
columns = [qbcolumns.Named(int, f"c{i}") for i in range(5)]
sel = dql.Select(columns)
proj_list = [0, 4, 2]
proj = sel.project(proj_list)
proj = sel.project(projlist)
expected_columns = qb.classes.stores.NamedStore(
((qbcolumns.Named(int, f"c{i}") for i in proj_list))
((qbcolumns.Named(int, n) for n in explist))
)
assert expected_columns == proj.column_origins
assert {} == proj.aliases
@pytest.mark.parametrize(
"projlist",
[
[0, 4, 0],
[0, 0],
["c0", "c4", "c0"],
["c0", "c0"],
[0, "c4", "c0"],
],
)
def test_project_only_columns_with_list(self, projlist):
columns = [qbcolumns.Named(int, f"c{i}") for i in range(5)]
sel = dql.Select(columns)
with pytest.raises(ValueError):
proj = sel.project(projlist)
def test_project_only_columns_with_slice(self):
columns = [qbcolumns.Named(int, f"c{i}") for i in range(5)]
sel = dql.Select(columns)
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment