Skip to content

Commit

Permalink
feat: add fields methods to QuerySet
Browse files Browse the repository at this point in the history
  • Loading branch information
jorwoods committed Feb 6, 2025
1 parent add9d18 commit f5afa19
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 0 deletions.
39 changes: 39 additions & 0 deletions tableauserverclient/server/endpoint/endpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
TypeVar,
Union,
)
from typing_extensions import Self

from tableauserverclient.models.pagination_item import PaginationItem
from tableauserverclient.server.request_options import RequestOptions
Expand Down Expand Up @@ -353,3 +354,41 @@ def paginate(self, **kwargs) -> QuerySet[T]:
@abc.abstractmethod
def get(self, request_options: Optional[RequestOptions] = None) -> tuple[list[T], PaginationItem]:
raise NotImplementedError(f".get has not been implemented for {self.__class__.__qualname__}")

def fields(self: Self, *fields: str) -> QuerySet:
"""
Add fields to the request options. If no fields are provided, the
default fields will be used. If fields are provided, the default fields
will be used in addition to the provided fields.
Parameters
----------
fields : str
The fields to include in the request options.
Returns
-------
QuerySet
"""
queryset = QuerySet(self)
queryset.request_options.fields |= set(fields) | set(("_default_",))
return queryset

def only_fields(self: Self, *fields: str) -> QuerySet:
"""
Add fields to the request options. If no fields are provided, the
default fields will be used. If fields are provided, the default fields
will be replaced by the provided fields.
Parameters
----------
fields : str
The fields to include in the request options.
Returns
-------
QuerySet
"""
queryset = QuerySet(self)
queryset.request_options.fields |= set(fields)
return queryset
36 changes: 36 additions & 0 deletions tableauserverclient/server/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,6 +208,42 @@ def paginate(self: Self, **kwargs) -> Self:
self.request_options.pagesize = kwargs["page_size"]
return self

def fields(self: Self, *fields: str) -> Self:
"""
Add fields to the request options. If no fields are provided, the
default fields will be used. If fields are provided, the default fields
will be used in addition to the provided fields.
Parameters
----------
fields : str
The fields to include in the request options.
Returns
-------
QuerySet
"""
self.request_options.fields |= set(fields) | set(("_default_"))
return self

def only_fields(self: Self, *fields: str) -> Self:
"""
Add fields to the request options. If no fields are provided, the
default fields will be used. If fields are provided, the default fields
will be replaced by the provided fields.
Parameters
----------
fields : str
The fields to include in the request options.
Returns
-------
QuerySet
"""
self.request_options.fields |= set(fields)
return self

@staticmethod
def _parse_shorthand_filter(key: str) -> tuple[str, str]:
tokens = key.split("__", 1)
Expand Down
10 changes: 10 additions & 0 deletions test/test_request_option.py
Original file line number Diff line number Diff line change
Expand Up @@ -368,3 +368,13 @@ def test_language_export(self) -> None:

resp = self.server.users.get_request(url, request_object=opts)
self.assertTrue(re.search("language=en-us", resp.request.query))

def test_queryset_fields(self) -> None:
loop = self.server.users.fields("id")
assert "id" in loop.request_options.fields
assert "_default_" in loop.request_options.fields

def test_queryset_only_fields(self) -> None:
loop = self.server.users.only_fields("id")
assert "id" in loop.request_options.fields
assert "_default_" not in loop.request_options.fields

0 comments on commit f5afa19

Please sign in to comment.