feat: add fields methods to QuerySet

Author: jorwoods

tableauserverclient/server/endpoint/endpoint.py

@@ -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
@@ -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

tableauserverclient/server/query.py

@@ -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)

test/test_request_option.py

@@ -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