diff --git a/.pylintrc b/.pylintrc index dd76acc..c262559 100644 --- a/.pylintrc +++ b/.pylintrc @@ -7,3 +7,4 @@ disable= too-many-arguments, too-many-instance-attributes, invalid-name, + too-many-lines, # __init__.py is > 1000 lines diff --git a/README.rst b/README.rst index a21b5d2..e399017 100644 --- a/README.rst +++ b/README.rst @@ -525,6 +525,49 @@ The attribute can be passed to the task payloads, in the ``attachment`` paramete ... ... ) + +Manage Teammates +________________ + +Manage the members of your Scale team via API. Check out `Scale Team API Documentation`__ for more information. + +__ https://docs.scale.com/reference/teams-overview + +List Teammates +^^^^^^^^^^^^^^ + +Lists all teammates in your Scale team. +Returns all teammates in a List of Teammate objects. + +.. code-block:: python + + teammates = client.list_teammates() + +Invite Teammate +^^^^^^^^^^^^^^^ + +Invites a list of email strings to your team with the provided role. +The available teammate roles are: 'labeler', 'member', or 'manager'. +Returns all teammates in a List of Teammate objects. + +.. code-block:: python + + from scaleapi import TeammateRole + + teammates = client.invite_teammates(['email1@example.com', 'email2@example.com'], TeammateRole.Member) + +Update Teammate Role +^^^^^^^^^^^^^^^^^^^^^ + +Updates a list of emails of your Scale team members with the new role. +The available teammate roles are: 'labeler', 'member', or 'manager'. +Returns all teammates in a List of Teammate objects. + +.. code-block python + + from scaleapi import TeammateRole + + teammates = client.update_teammates_role(['email1@example.com', 'email2@example.com'], TeammateRole.Manager) Example Scripts _______________ @@ -612,6 +655,139 @@ Create a training task. client.create_training_task(TaskType, ...task parameters...) +Studio Assignments (For Scale Studio only) +__________________________________________ + +Manage project assignments for your labelers. + +List All Assignments +^^^^^^^^^^^^^^^^^^^^ + +Lists all your Scale team members and the projects they are assigned to. +Returns a dictionary of all teammate assignments with keys as 'emails' of each teammate, and values as a list of project names the teammate are assigned to. + +.. code-block:: python + + assignments = client.list_studio_assignments() + my_assignment = assignments.get('my-email@example.com') + +Add Studio Assignment +^^^^^^^^^^^^^^^^^^^^^ + +Assigns provided projects to specified teammate emails. + +Accepts a list of emails and a list of projects. + +Returns a dictionary of all teammate assignments with keys as 'emails' of each teammate, and values as a list of project names the teammate are assigned to. + +.. code-block:: python + + assignments = client.add_studio_assignments(['email1@example.com', 'email2@example.com'], ['project 1', 'project 2']) + + +Remove Studio Assignment +^^^^^^^^^^^^^^^^^^^^^^^^ + +Removes provided projects from specified teammate emails. + +Accepts a list of emails and a list of projects. + +Returns a dictionary of all teammate assignments with keys as 'emails' of each teammate, and values as a list of project names the teammate are assigned to. + +.. code-block:: python + + assignments = client.remove_studio_assignments(['email1@example.com', 'email2@example.com'], ['project 1', 'project 2']) + +Studio Project Groups (For Scale Studio Only) +_____________________________________________ + +Manage groups of labelers in our project by using Studio Project Groups. + +List Studio Project Groups +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Returns all labeler groups for the specified project. + +.. code-block:: python + + list_project_group = client.list_project_groups('project_name') + +Add Studio Project Group +^^^^^^^^^^^^^^^^^^^^^^^^ + +Creates a project group with the provided group_name for the specified project and adds the provided teammate emails to the new project group. The team members must be assigned to the specified project in order to be added to the new group. + +Returns the created StudioProjectGroup object. + +.. code-block:: python + + added_project_group = client.create_project_group( + 'project_name', ['email1@example.com'], 'project_group_name' + ) + +Update Studio Project Group +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Assign or remove teammates from a project group. + +Returns the updated StudioProjectGroup object. + +.. code-block:: python + + updated_project_group = client.update_project_group( + 'project_name', 'project_group_name', ['emails_to_add'], ['emails_to_remove'] + ) + +Studio Batches (For Scale Studio Only) +_______________________________________ + +Get information about your pending Studio batches. + +List Studio Batches +^^^^^^^^^^^^^^^^^^^ + +Returns a list of StudioBatch objects for all pending Studio batches. + +.. code-block:: python + + studio_batches = client.list_studio_batches() + +Assign Studio Batches +^^^^^^^^^^^^^^^^^^^^^^ + +Sets labeler group assignment for the specified batch. + +Returns a StudioBatch object for the specified batch. + +.. code-block:: python + + assigned_studio_batch = client.assign_studio_batches('batch_name', ['project_group_name']) + +Set Studio Batches Priority +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Sets the order to prioritize your pending Studio batches. You must include all pending studio batches in the List. + +Returns a List of StudioBatch objects in the new order. + +.. code-block:: python + + studio_batch_priority = client.set_studio_batches_priorities( + ['pending_batch_1', 'pending_batch_2', 'pending_batch_3'] + ) + +Reset Studio Batches Priority +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +Resets the order of your Studio batches to the default order, which prioritizes older batches first. + +Returns a List of StudioBatch objects in the new order. + +.. code-block:: python + + reset_studio_batch_prioprity = client.reset_studio_batches_priorities() + + Error handling ______________ @@ -642,7 +818,9 @@ For example: print(err.code) # 400 print(err.message) # Parameter is invalid, reason: "attachments" is required + + Troubleshooting _______________ -If you notice any problems, please email us at support@scale.com. +If you notice any problems, please contact our support via Intercom by logging into your dashboard, or, if you are Enterprise, by contacting your Engagement Manager. diff --git a/scaleapi/__init__.py b/scaleapi/__init__.py index e6bb922..0e3b9b5 100644 --- a/scaleapi/__init__.py +++ b/scaleapi/__init__.py @@ -9,7 +9,9 @@ from ._version import __version__ # noqa: F401 from .api import Api +from .studio import StudioBatch, StudioLabelerAssignment, StudioProjectGroup from .tasks import Task, TaskReviewStatus, TaskStatus, TaskType +from .teams import Teammate, TeammateRole T = TypeVar("T") @@ -918,3 +920,234 @@ def create_training_task( training_task_data = self.api.post_request(endpoint, body=kwargs) return TrainingTask(training_task_data, self) + + def list_teammates(self) -> List[Teammate]: + """Returns all teammates. + Refer to Teams API Reference: + https://docs.scale.com/reference/teams-list + Returns: + List[Teammate] + """ + endpoint = "teams" + teammate_list = self.api.get_request(endpoint) + return [Teammate(teammate, self) for teammate in teammate_list] + + def invite_teammates(self, emails: List[str], role: TeammateRole) -> List[Teammate]: + """Invites a list of emails to your team. + + Args: + emails (List[str]): + emails to invite + role (TeammateRole): + role to invite + Returns: + List[Teammate] + """ + endpoint = "teams/invite" + payload = { + "emails": emails, + "team_role": role.value, + } + teammate_list = self.api.post_request(endpoint, payload) + return [Teammate(teammate, self) for teammate in teammate_list] + + def update_teammates_role( + self, emails: List[str], role: TeammateRole + ) -> List[Teammate]: + """Updates role of teammates by email + + Args: + emails (List[str]): + emails to update + role (TeammateRole): + new role + Returns: + List[Teammate] + """ + endpoint = "teams/set_role" + payload = { + "emails": emails, + "team_role": role.value, + } + teammate_list = self.api.post_request(endpoint, payload) + return [Teammate(teammate, self) for teammate in teammate_list] + + def list_studio_assignments(self) -> Dict[str, StudioLabelerAssignment]: + """Returns a dictionary where the keys are user emails and the + values are projects the user is assigned to. + + Returns: + Dict[StudioLabelerAssignment] + """ + endpoint = "studio/assignments" + raw_assignments = self.api.get_request(endpoint) + assignments = {} + for (email, assigned_projects) in raw_assignments.items(): + assignments[email] = StudioLabelerAssignment(assigned_projects, email, self) + return assignments + + def add_studio_assignments( + self, emails: List[str], projects: List[str] + ) -> Dict[str, StudioLabelerAssignment]: + """Adds projects to the users based on emails. + Args: + emails (List[str]): + emails to assign + projects (List[str]): + projects to assign + + Returns: + Dict[StudioLabelerAssignment] + """ + endpoint = "studio/assignments/add" + payload = { + "emails": emails, + "projects": projects, + } + raw_assignments = self.api.post_request(endpoint, payload) + assignments = {} + for (email, assigned_projects) in raw_assignments.items(): + assignments[email] = StudioLabelerAssignment(assigned_projects, email, self) + return assignments + + def remove_studio_assignments( + self, emails: List[str], projects: List[str] + ) -> Dict[str, StudioLabelerAssignment]: + """Removes projects from users based on emails. + Args: + emails (List[str]): + emails to unassign + projects (List[str]): + projects to unassign + + Returns: + Dict[StudioLabelerAssignment] + """ + endpoint = "studio/assignments/remove" + payload = { + "emails": emails, + "projects": projects, + } + raw_assignments = self.api.post_request(endpoint, payload) + assignments = {} + for (email, assigned_projects) in raw_assignments.items(): + assignments[email] = StudioLabelerAssignment(assigned_projects, email, self) + return assignments + + def list_project_groups(self, project: str) -> List[StudioProjectGroup]: + """List all labeler groups for the specified project. + Args: + project (str): + project to retrieve labeler groups from + + Returns: + List[StudioProjectGroup] + """ + endpoint = f"studio/projects/{Api.quote_string(project)}/groups" + groups = self.api.get_request(endpoint) + return [StudioProjectGroup(group, self) for group in groups] + + # StudioWorker for each worker in a group + + def create_project_group( + self, project: str, emails: List[str], project_group: str + ) -> StudioProjectGroup: + """Creates a labeler group for the specified project. + Args: + project (str): + project to create a labeler group in + emails (List[str]): + list of labeler emails to add to the project group + project_group (str): + name of the project group to create + + Returns: + StudioProjectGroup + """ + endpoint = f"studio/projects/{Api.quote_string(project)}/groups" + payload = {"emails": emails, "name": project_group} + return StudioProjectGroup(self.api.post_request(endpoint, payload), self) + + def update_project_group( + self, + project: str, + project_group: str, + add_emails: List[str], + remove_emails: List[str], + ) -> StudioProjectGroup: + """Updates specified labeler group for the specified project. + Args: + project (str): + project to create a labeler group in + project_group (str): + name of the project group to create + add_emails (List[str]): + list of labeler emails to add to the project group + remove_emails (List[str]): + list of labeler emails to remove to the project group + + Returns: + StudioProjectGroup + """ + endpoint = ( + f"studio/projects/{Api.quote_string(project)}" + f"/groups/{Api.quote_string(project_group)}" + ) + payload = {"add_emails": add_emails, "remove_emails": remove_emails} + return StudioProjectGroup(self.api.put_request(endpoint, payload), self) + + def list_studio_batches(self) -> List[StudioBatch]: + """Returns a list with all pending studio batches, + in order of priority. + + Returns: + List[StudioBatch] + """ + endpoint = "studio/batches" + batches = self.api.get_request(endpoint) + return [StudioBatch(batch, self) for batch in batches] + + # StudioBatchStatus for each batch_type in a batch + + def assign_studio_batches( + self, batch_name: str, project_groups: List[str] + ) -> StudioBatch: + """Sets labeler group assignment for the specified batch. + Args: + batch_name (str): + batch name to assign project_groups to + project_groups (List[str]): + project groups to be assigned + Returns: + StudioBatch + """ + endpoint = f"studio/batches/{Api.quote_string(batch_name)}" + payload = {"groups": project_groups} + return StudioBatch(self.api.put_request(endpoint, payload), self) + + def set_studio_batches_priorities( + self, batch_names: List[str] + ) -> List[StudioBatch]: + """Sets the priority of batches based on the array order. + Args: + batches (List[str]): + list of all pending batches names ordered by priority + Returns: + List[StudioBatch] + """ + batches_names = list(map(lambda batch_name: {"name": batch_name}, batch_names)) + endpoint = "studio/batches/set_priorities" + payload = {"batches": batches_names} + batches = self.api.post_request(endpoint, payload) + return [StudioBatch(batch, self) for batch in batches] + + def reset_studio_batches_priorities(self) -> List[StudioBatch]: + """Resets the priority of batches. (Default order is + sorted by creation date) + + Returns: + List[StudioBatch] + """ + endpoint = "studio/batches/reset_priorities" + batches = self.api.post_request(endpoint) + return [StudioBatch(batch, self) for batch in batches] diff --git a/scaleapi/_version.py b/scaleapi/_version.py index bcc5c8a..e2af5ab 100644 --- a/scaleapi/_version.py +++ b/scaleapi/_version.py @@ -1,2 +1,2 @@ -__version__ = "2.10.1" +__version__ = "2.11.0" __package_name__ = "scaleapi" diff --git a/scaleapi/api.py b/scaleapi/api.py index 078f305..05d89cf 100644 --- a/scaleapi/api.py +++ b/scaleapi/api.py @@ -159,6 +159,17 @@ def delete_request(self, endpoint, params=None): "DELETE", endpoint, headers=self._headers, auth=self._auth, params=params ) + def put_request(self, endpoint, body=None, params=None): + """Generic PUT Request Wrapper""" + return self._api_request( + "PUT", + endpoint, + body=body, + headers=self._headers, + auth=self._auth, + params=params, + ) + @staticmethod def _generate_useragent(extension: str = None) -> str: """Generates UserAgent parameter with module, Python diff --git a/scaleapi/studio.py b/scaleapi/studio.py new file mode 100644 index 0000000..cb663aa --- /dev/null +++ b/scaleapi/studio.py @@ -0,0 +1,116 @@ +from enum import Enum + + +class StudioLabelerAssignment: + """Labeler Assignment class, contains information about + assignments for labelers.""" + + def __init__(self, assigned_projects, email, client): + self._json = assigned_projects + self._client = client + self.email = email + self.assigned_projects = assigned_projects + + def __hash__(self): + return hash(self.email) + + def __str__(self): + return ( + f"StudioLabelerAssignment(email={self.email}," + f"assigned_projects={self.assigned_projects})" + ) + + def __repr__(self): + return f"StudioLabelerAssignment({self._json})" + + def as_dict(self): + """Returns all attributes as a dictionary""" + return self._json + + +class StudioWorker: + """Worker object that is returned in the 'workers' array.""" + + def __init__(self, json, client): + self._json = json + self._client = client + self.id = json["id"] + self.email = json["email"] + + def __hash__(self): + return hash(self.id) + + def __str__(self): + return f"StudioWorker(id={self.email})" + + def __repr__(self): + return f"StudioWorker({self._json})" + + def as_dict(self): + """Returns all attributes as a dictionary""" + return self._json + + +class StudioProjectGroup: + """Studio project group.""" + + def __init__(self, json, client): + self._json = json + self._client = client + self.id = json["id"] + self.name = json["name"] + self.numWorkers = json["numWorkers"] + self.isSingleton = json["isSingleton"] + if json.get("workers"): + self.workers = ([StudioWorker(w, client) for w in json["workers"]],) + else: + self.workers = [] + + def __hash__(self): + return hash(self.id) + + def __str__(self): + return f"StudioProjectGroup(name={self.name})" + + def __repr__(self): + return f"StudioProjectGroup({self._json})" + + def as_dict(self): + """Returns all attributes as a dictionary""" + return self._json + + +class StudioBatchStatus(Enum): + """Studio Batch Statuses""" + + Production = "Production" + + +class StudioBatch: + """Studio Batch""" + + def __init__(self, json, client): + self._json = json + self._client = client + self.id = json["id"] + self.name = json["name"] + self.project_id = json["projectId"] + self.project_name = json["projectName"] + self.batch_type = json["batchType"] + self.studio_priority = json.get("studioPriority") + self.total = json["total"] + self.completed = json["completed"] + self.groups = json["groups"] + + def __hash__(self): + return hash(self.id) + + def __str__(self): + return f"StudioBatch(name={self.name})" + + def __repr__(self): + return f"StudioBatch({self._json})" + + def as_dict(self): + """Returns all attributes as a dictionary""" + return self._json diff --git a/scaleapi/teams.py b/scaleapi/teams.py new file mode 100644 index 0000000..9f05bc6 --- /dev/null +++ b/scaleapi/teams.py @@ -0,0 +1,43 @@ +from enum import Enum + + +class TeammateRole(Enum): + """Teammate Roles Enum""" + + Labeler = "labeler" + Member = "member" + Manager = "manager" + Admin = "admin" + + +class Teammate: + """Teammate class, containing teammate information.""" + + def __init__(self, json, client): + self._json = json + self.email: str = json["email"] + self.role = json["role"] + self._client = client + # fill in rest here (non-optional fields) + self.company = json.get("company") + self.first_name = json.get("firstName") + self.last_name = json.get("lastName") + self.is_studio_labeler = json.get("isStudioLabeler") + self.expiry = json.get("expiry") + + def __hash__(self): + return hash(self.email) + + def __str__(self): + return f"Teammate(email={self.email})" + + def __repr__(self): + return f"Teammate({self._json})" + + def as_dict(self): + """Returns all attributes as a dictionary""" + return self._json + + def update_teammate_role(self, new_role: TeammateRole): + """Updates teammate role""" + return self._client.update_teammates_role([self.email], new_role) diff --git a/tests/t_studio.py b/tests/t_studio.py new file mode 100644 index 0000000..9f25189 --- /dev/null +++ b/tests/t_studio.py @@ -0,0 +1,103 @@ +import os +import time +import uuid + +import scaleapi + +# Testing with live key since Studio endpoints require live projects + +client = scaleapi.ScaleClient(os.environ["JF_STUDIO_KEY"]) + +test_project = "vfps" +# test teammates + +# test list_teammates +teammates = client.list_teammates() +for teammate in teammates: + print("Testing list_teammates:") + print(teammate.as_dict()) + break + +# test invite_teammate +test_email_invite = f"jon.feng+sdk_test_{str(uuid.uuid4())[-4:]}@scale.com" +new_teammates = client.invite_teammates( + [test_email_invite], scaleapi.TeammateRole.Member +) + +for new_teammate in new_teammates: + if new_teammate.email == test_email_invite: + print("testing invite teammate:") + print(new_teammate.as_dict()) + +time.sleep(20) + +updated_teammates = client.update_teammates_role( + [test_email_invite], scaleapi.TeammateRole.Manager +) + +for updated_teammate in updated_teammates: + if updated_teammate.email == test_email_invite: + print("testing update teammate:") + print(updated_teammate.as_dict()) + +# test assignments +assignments = client.list_studio_assignments() +print("Testing listing assignments:") +print(assignments["jon.feng@scale.com"]) + +added_assignment = client.add_studio_assignments([test_email_invite], [test_project]) + +print("Testing adding assignment:") +print(added_assignment.get(test_email_invite)) + +removed_assignment = client.remove_studio_assignments( + [test_email_invite], [test_project] +) +print("Testing removing assignment:") +print(removed_assignment.get(test_email_invite)) +# re-add assignment for next step +added_assignment = client.add_studio_assignments([test_email_invite], [test_project]) + +project_group_name = "sdk-testing" +added_project_group = client.create_project_group( + test_project, [test_email_invite], project_group_name +) + +print("Test creating project group:") +print(added_project_group.as_dict()) + +list_project_group = client.list_project_groups(test_project) +print("Test listing project groups") +print(list_project_group) + +updated_project_group = client.update_project_group( + test_project, project_group_name, [], [test_email_invite] +) +print("Test removing project group") +print(updated_project_group.as_dict()) + +re_updated_project_group = client.update_project_group( + test_project, project_group_name, [test_email_invite], [] +) +print("Test adding back project group") +print(re_updated_project_group.as_dict()) + +studio_batches = client.list_studio_batches() +print("Test studio batches") +print(studio_batches) + +test_batch = "testing_vfps" + +assigned_studio_batch = client.assign_studio_batches(test_batch, [project_group_name]) +print("Test studio batch assignment") +print(assigned_studio_batch) + +studio_batch_priority = client.set_studio_batches_priorities( + list(map(lambda sb: sb.name, studio_batches)) +) +print("Test set studio batch priority") +print(studio_batch_priority) + +reset_studio_batch_prioprity = client.reset_studio_batches_priorities() +print("Test reset studio batch priority") +print(reset_studio_batch_prioprity) diff --git a/tests/test_client.py b/tests/test_client.py index ed3ddc1..65ce105 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -1,6 +1,5 @@ # pylint: disable=missing-function-docstring -import os import time import uuid from datetime import datetime @@ -16,12 +15,13 @@ ScaleUnauthorized, ) from scaleapi.tasks import TaskType +from scaleapi.teams import TeammateRole TEST_PROJECT_NAME = "scaleapi-python-sdk" try: print(f"SDK Version: {scaleapi.__version__}") - test_api_key = os.environ["SCALE_TEST_API_KEY"] + test_api_key = "test_fe79860cdbe547bf91b4e7da897a6c92" if test_api_key.startswith("test_") or test_api_key.endswith("|test"): client = scaleapi.ScaleClient(test_api_key, "pytest") @@ -435,3 +435,40 @@ def test_files_import(): file_url="https://static.scale.com/uploads/selfserve-sample-image.png", project_name=TEST_PROJECT_NAME, ) + + +current_timestamp = str(uuid.uuid4)[-9:] +TEST_USER = f"test+{current_timestamp}@scale.com" + + +def test_list_teammates(): + teammates = client.list_teammates() + assert len(teammates) > 0 + + +def test_invite_teammates(): + old_teammates = client.list_teammates() + new_teammates = client.invite_teammates([TEST_USER], TeammateRole.Member) + assert len(new_teammates) >= len( + old_teammates + ) # needs to sleep for teammates list to be updated + + +STUDIO_TEST_PROJECT = "python-sdk-studio-test" + +try: + project = client.get_project(STUDIO_TEST_PROJECT) +except ScaleResourceNotFound: + client.create_project(project_name=STUDIO_TEST_PROJECT) +STUDIO_BATCH_TEST_NAME = f"studio-test-batch-{current_timestamp}" + + +def test_list_studio_batches(): + # Create a test project if it does not already exist + + # Create a test batch + client.create_batch(STUDIO_TEST_PROJECT, STUDIO_BATCH_TEST_NAME) + + # Check that the batch is returned by the list_studio_batches method + batches = client.list_studio_batches() + assert len(batches) > 0