Add training tasks endpoint to Python client (#50)

* Add training tasks endpoint to Python client

* fix

* lint

Author: Tony Wu

README.rst

@@ -528,6 +528,19 @@ __ https://docs.scale.com/reference
 
     client.create_evaluation_task(TaskType.TextCollection, **payload)
 
+Training tasks (For Scale Rapid projects only)
+________________________________________________
+
+Training tasks are used to onboard taskers onto your project
+
+Create Training Task
+^^^^^^^^^^^^^^^^^^^^^^
+
+Create a training task.
+
+.. code-block:: python
+
+    client.create_training_task(TaskType, ...task parameters...)
 
 Error handling
 ______________

scaleapi/__init__.py

@@ -5,6 +5,7 @@
 from scaleapi.exceptions import ScaleInvalidRequest
 from scaleapi.files import File
 from scaleapi.projects import Project
+from scaleapi.training_tasks import TrainingTask
 
 from ._version import __version__  # noqa: F401
 from .api import Api
@@ -794,10 +795,12 @@ def create_evaluation_task(
         task_type: TaskType,
         **kwargs,
     ) -> EvaluationTask:
-        """This method can only be used for Self-Serve projects.
+        """This method can only be used for Rapid projects.
         Supported Task Types: [
+            DocumentTranscription,
+            SegmentAnnotation,
+            VideoPlaybackAnnotation,
             ImageAnnotation,
-            Categorization,
             TextCollection,
             NamedEntityRecognition
         ]
@@ -827,3 +830,38 @@ def create_evaluation_task(
 
         evaluation_task_data = self.api.post_request(endpoint, body=kwargs)
         return EvaluationTask(evaluation_task_data, self)
+
+    def create_training_task(
+        self,
+        task_type: TaskType,
+        **kwargs,
+    ) -> TrainingTask:
+        """This method can only be used for Rapid projects.
+        Supported Task Types: [
+            DocumentTranscription,
+            SegmentAnnotation,
+            VideoPlaybackAnnotation,
+            ImageAnnotation,
+            TextCollection,
+            NamedEntityRecognition
+        ]
+        Parameters may differ based on the given task_type.
+
+        Args:
+            task_type (TaskType):
+                Task type to be created
+                e.g.. `TaskType.ImageAnnotation`
+            **kwargs:
+                The same set of parameters are expected with
+                create_task function and an additional expected_response.
+                Scale's API documentation.
+                https://docs.scale.com/reference
+
+        Returns:
+            TrainingTask:
+                Returns created training task.
+        """
+        endpoint = f"training_tasks/{task_type.value}"
+
+        training_task_data = self.api.post_request(endpoint, body=kwargs)
+        return TrainingTask(training_task_data, self)

scaleapi/training_tasks.py

@@ -0,0 +1,22 @@
+class TrainingTask:
+    """TrainingTask class, containing TrainingTask information."""
+
+    def __init__(self, json, client):
+        self._json = json
+        self.id = json["id"]
+        self.initial_response = getattr(json, "initial_response", None)
+        self.expected_response = json["expected_response"]
+        self._client = client
+
+    def __hash__(self):
+        return hash(self.id)
+
+    def __str__(self):
+        return f"TrainingTask(id={self.id})"
+
+    def __repr__(self):
+        return f"TrainingTask({self._json})"
+
+    def as_dict(self):
+        """Returns all attributes as a dictionary"""
+        return self._json