From d5e8510faa2922912f7400bf0d65e7884a67886d Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Mon, 19 Dec 2022 14:15:05 +0000 Subject: [PATCH] minor tweak to `M2M` so it flattens `all_columns` --- piccolo/columns/m2m.py | 20 +++++++++++++------- piccolo/query/mixins.py | 9 ++------- piccolo/utils/list.py | 27 +++++++++++++++++++++++++++ tests/columns/test_m2m.py | 35 +++++++++++++++++++++++++++++++++++ tests/utils/test_list.py | 8 ++++++++ 5 files changed, 85 insertions(+), 14 deletions(-) create mode 100644 piccolo/utils/list.py create mode 100644 tests/utils/test_list.py diff --git a/piccolo/columns/m2m.py b/piccolo/columns/m2m.py index 5c696bd2b..bb7f30b39 100644 --- a/piccolo/columns/m2m.py +++ b/piccolo/columns/m2m.py @@ -12,6 +12,7 @@ ForeignKey, LazyTableReference, ) +from piccolo.utils.list import flatten from piccolo.utils.sync import run_sync if t.TYPE_CHECKING: @@ -45,9 +46,9 @@ def __init__( self.m2m = m2m self.load_json = load_json - safe_types = [int, str] + safe_types = (int, str) - # If the columns can be serialised / deserialise as JSON, then we + # If the columns can be serialised / deserialised as JSON, then we # can fetch the data all in one go. self.serialisation_safe = all( (column.__class__.value_type in safe_types) @@ -408,7 +409,10 @@ def __init__( ) def __call__( - self, *columns: Column, as_list: bool = False, load_json: bool = False + self, + *columns: t.Union[Column, t.List[Column]], + as_list: bool = False, + load_json: bool = False, ) -> M2MSelect: """ :param columns: @@ -420,14 +424,16 @@ def __call__( :param load_json: If ``True``, any JSON strings are loaded as Python objects. """ - if not columns: - columns = tuple(self._meta.secondary_table._meta.columns) + columns_ = flatten(columns) + + if not columns_: + columns_ = self._meta.secondary_table._meta.columns - if as_list and len(columns) != 1: + if as_list and len(columns_) != 1: raise ValueError( "`as_list` is only valid with a single column argument" ) return M2MSelect( - *columns, m2m=self, as_list=as_list, load_json=load_json + *columns_, m2m=self, as_list=as_list, load_json=load_json ) diff --git a/piccolo/query/mixins.py b/piccolo/query/mixins.py index 36a2fd80f..582c8bbb1 100644 --- a/piccolo/query/mixins.py +++ b/piccolo/query/mixins.py @@ -10,6 +10,7 @@ from piccolo.columns.column_types import ForeignKey from piccolo.custom_types import Combinable from piccolo.querystring import QueryString +from piccolo.utils.list import flatten from piccolo.utils.sql_values import convert_to_sql_value if t.TYPE_CHECKING: # pragma: no cover @@ -440,13 +441,7 @@ def columns(self, *columns: t.Union[Selectable, t.List[Selectable]]): in which case we flatten the list. """ - _columns = [] - for column in columns: - if isinstance(column, list): - _columns.extend(column) - else: - _columns.append(column) - + _columns = flatten(columns) combined = list(self.selected_columns) + _columns self.selected_columns = combined diff --git a/piccolo/utils/list.py b/piccolo/utils/list.py new file mode 100644 index 000000000..b26847adc --- /dev/null +++ b/piccolo/utils/list.py @@ -0,0 +1,27 @@ +import typing as t + +ElementType = t.TypeVar("ElementType") + + +def flatten( + items: t.Sequence[t.Union[ElementType, t.List[ElementType]]] +) -> t.List[ElementType]: + """ + Takes a sequence of elements, and flattens it out. For example:: + + >>> flatten(['a', ['b', 'c']]) + ['a', 'b', 'c'] + + We need this for situations like this:: + + await Band.select(Band.name, Band.manager.all_columns()) + + """ + _items: t.List[ElementType] = [] + for item in items: + if isinstance(item, list): + _items.extend(item) + else: + _items.append(item) + + return _items diff --git a/tests/columns/test_m2m.py b/tests/columns/test_m2m.py index 4dff73430..653202ed8 100644 --- a/tests/columns/test_m2m.py +++ b/tests/columns/test_m2m.py @@ -273,6 +273,41 @@ def test_select_id(self): ], ) + @engines_skip("cockroach") + def test_select_all_columns(self): + """ + Make sure ``all_columns`` can be passed in as an argument. ``M2M`` + should flatten the arguments. Reported here: + + https://github.com/piccolo-orm/piccolo/issues/728 + + 🐛 Cockroach bug: https://github.com/cockroachdb/cockroach/issues/71908 "could not decorrelate subquery" error under asyncpg + + """ # noqa: E501 + response = Band.select( + Band.name, Band.genres(Genre.all_columns(exclude=(Genre.id,))) + ).run_sync() + self.assertEqual( + response, + [ + { + "name": "Pythonistas", + "genres": [ + {"name": "Rock"}, + {"name": "Folk"}, + ], + }, + {"name": "Rustaceans", "genres": [{"name": "Folk"}]}, + { + "name": "C-Sharps", + "genres": [ + {"name": "Rock"}, + {"name": "Classical"}, + ], + }, + ], + ) + def test_add_m2m(self): """ Make sure we can add items to the joining table. diff --git a/tests/utils/test_list.py b/tests/utils/test_list.py new file mode 100644 index 000000000..e0211c6f7 --- /dev/null +++ b/tests/utils/test_list.py @@ -0,0 +1,8 @@ +from unittest import TestCase + +from piccolo.utils.list import flatten + + +class TestFlatten(TestCase): + def test_flatten(self): + self.assertListEqual(flatten(["a", ["b", "c"]]), ["a", "b", "c"])