Skip to content

Commit

Permalink
allow multiple keys in ui.get
Browse files Browse the repository at this point in the history
  • Loading branch information
rodja committed Nov 1, 2023
1 parent 5735aa5 commit 284ccba
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 9 deletions.
22 changes: 14 additions & 8 deletions nicegui/get.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
from typing import Generic, Iterator, List, Optional, Type, TypeVar
from typing import Generic, Iterator, List, Optional, Type, TypeVar, Union

from typing_extensions import Self

Expand All @@ -10,20 +10,26 @@

class elements(Generic[T], Iterator[T]):

def __init__(self, *, type: Optional[Type[T]] = Element, key: str = '') -> None:
self.type = type
self.key = key
self._within_types = []
self._within_keys = []
def __init__(self, *, type: Type[T] = Element, key: Union[str, list[str], None] = None) -> None:
"""Get elements by type and/or key.
:param type: type of the elements to get (iterator will then be of type `type`)
:param key: key of the elements to get, can be a list of strings or a single string where keys are separated by whitespace
"""
self._types = type
self._keys = key.split() if isinstance(key, str) else key
assert self._keys is None or isinstance(self._keys, list)
self._within_types: list[Element] = []
self._within_keys: list[str] = []

def __iter__(self) -> Iterator[T]:
client = context.get_client()
return self.iterate(client.layout)

def iterate(self, parent: Element, *, visited: List[Element] = []) -> Iterator[T]:
for element in parent:
if (self.type is None or isinstance(element, self.type)) and \
(not self.key or self.key in element._keys):
if (self._types is None or isinstance(element, self._types)) and \
(not self._keys or all(key in element._keys for key in self._keys)):
if (not self._within_types or any(isinstance(element, type) for type in self._within_types for element in visited)) and \
(not self._within_keys or any(key in element._keys for key in self._within_keys for element in visited)):
yield element
Expand Down
14 changes: 13 additions & 1 deletion tests/test_get_elements.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,7 +80,7 @@ def test_get_by_key(screen: Screen):
assert result[0].text == 'button B'


def test_get_by_multiple_keys(screen: Screen):
def test_get_by_specific_key(screen: Screen):
ui.button('button A').keys('test')
ui.button('button B').keys('important ', 'test')
ui.button('button C').keys(' important test')
Expand All @@ -98,6 +98,18 @@ def test_get_by_multiple_keys(screen: Screen):
assert important[1].text == 'button C'


def test_get_by_multiple_keys(screen: Screen):
ui.button('button A').keys('test')
ui.button('button B').keys('important ', 'test')
ui.button('button C').keys(' important test')

search = ui.get(type=ui.button, key='test important')
result = [b.text for b in search]
screen.open('/')
ic(search.__dict__)
assert result == ['button B', 'button C']


def test_get_within_type(screen: Screen):
ui.button('button A')
ui.label('label A')
Expand Down

0 comments on commit 284ccba

Please sign in to comment.