use set in user._gather_elements and improve typing

zauberzeug · Jul 18, 2024 · 101d5bf · 101d5bf
1 parent d6c31c7
commit 101d5bf
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 15 deletions.
diff --git a/nicegui/testing/user.py b/nicegui/testing/user.py
@@ -2,7 +2,7 @@
 
 import asyncio
 import re
-from typing import List, Optional, Type, TypeVar, Union, overload
+from typing import List, Optional, Set, Type, TypeVar, Union, overload
 from uuid import uuid4
 
 import httpx
@@ -74,7 +74,7 @@ async def should_see(self,
     @overload
     async def should_see(self,
                          *,
-                         kind: Type[T] = Element,
+                         kind: Optional[Type[T]] = None,
                          marker: Union[str, list[str], None] = None,
                          content: Union[str, list[str], None] = None,
                          retries: int = 3,
@@ -84,7 +84,7 @@ async def should_see(self,
     async def should_see(self,
                          target: Union[str, Type[T], None] = None,
                          *,
-                         kind: Type[T] = Element,
+                         kind: Optional[Type[T]] = None,
                          marker: Union[str, list[str], None] = None,
                          content: Union[str, list[str], None] = None,
                          retries: int = 3,
@@ -114,7 +114,7 @@ async def should_not_see(self,
     @overload
     async def should_not_see(self,
                              *,
-                             kind: Type[T] = Element,
+                             kind: Optional[Type[T]] = None,
                              marker: Union[str, list[str], None] = None,
                              content: Union[str, list[str], None] = None,
                              retries: int = 3,
@@ -124,7 +124,7 @@ async def should_not_see(self,
     async def should_not_see(self,
                              target: Union[str, Type[T], None] = None,
                              *,
-                             kind: Type[T] = Element,
+                             kind: Optional[Type[T]] = None,
                              marker: Union[str, list[str], None] = None,
                              content: Union[str, list[str], None] = None,
                              retries: int = 3,
@@ -156,7 +156,7 @@ def find(self,
     def find(self,
              target: Union[str, Type[T], None] = None,
              *,
-             kind: Type[T] = Element,
+             kind: Optional[Type[T]] = None,
              marker: Union[str, list[str], None] = None,
              content: Union[str, list[str], None] = None,
              ) -> UserInteraction:
@@ -177,29 +177,33 @@ def current_layout(self) -> Element:
 
     def _gather_elements(self,
                          target: Union[str, Type[T], None] = None,
-                         kind: Type[T] = Element,
+                         kind: Optional[Type[T]] = None,
                          marker: Union[str, list[str], None] = None,
                          content: Union[str, list[str], None] = None,
-                         ) -> List[T]:
+                         ) -> Set[T]:
         if target is None:
-            return list(ElementFilter(kind=kind, marker=marker, content=content))
+            return set(ElementFilter(kind=kind, marker=marker, content=content))
         elif isinstance(target, str):
-            return list(list(ElementFilter(marker=target)) + list(ElementFilter(content=target)))
+            elements: Set[T] = set(ElementFilter(marker=target))
+            elements.update(ElementFilter(content=target))
+            return elements
         else:
-            return list(ElementFilter(kind=target))
+            return set(ElementFilter(kind=target))
 
     def _build_error_message(self,
                              target: Union[str, Type[T], None] = None,
-                             kind: Type[T] = Element,
+                             kind: Optional[Type[T]] = None,
                              marker: Union[str, list[str], None] = None,
                              content: Union[str, list[str], None] = None,
                              ) -> str:
         if isinstance(target, str):
             return f'element with marker={target} or content={target} on the page:\n{self.current_layout}'
         elif target is not None:
             return f'element of type {target.__name__} on the page:\n{self.current_layout}'
-        else:
+        elif kind is not None:
             return f'element of type {kind.__name__} with {marker=} and {content=} on the page:\n{self.current_layout}'
+        else:
+            return f'element with {marker=} and {content=} on the page:\n{self.current_layout}'
 
 
 original_get_slot_stack = Slot.get_stack

diff --git a/nicegui/testing/user_interaction.py b/nicegui/testing/user_interaction.py
@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, List, TypeVar
+from typing import TYPE_CHECKING, Set, TypeVar
 
 from typing_extensions import Self
 
@@ -15,7 +15,7 @@
 
 class UserInteraction:
 
-    def __init__(self, user: User, elements: List[T]) -> None:
+    def __init__(self, user: User, elements: Set[T]) -> None:
         """Iteraction object of the simulated user.
 
         This will be returned by the ``find`` method of the ``user`` fixture in pytests.