From 6abddc051ecdf1f971dcd68c493b28e9b3129986 Mon Sep 17 00:00:00 2001
From: Mark Daoust <markdaoust@google.com>
Date: Fri, 28 Jun 2024 14:59:47 -0700
Subject: [PATCH] Add tuned_models samples

Change-Id: I4b85a0179d2e7e4c112ed6250a7e602de57787c4
---
 google/generativeai/types/__init__.py         |  1 +
 google/generativeai/types/permission_types.py | 70 ++++++++++++++++++-
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/google/generativeai/types/__init__.py b/google/generativeai/types/__init__.py
index 21768bbe6..0acfb1397 100644
--- a/google/generativeai/types/__init__.py
+++ b/google/generativeai/types/__init__.py
@@ -21,6 +21,7 @@
 from google.generativeai.types.generation_types import *
 from google.generativeai.types.helper_types import *
 from google.generativeai.types.model_types import *
+from google.generativeai.types.permission_types import *
 from google.generativeai.types.safety_types import *
 from google.generativeai.types.text_types import *
 
diff --git a/google/generativeai/types/permission_types.py b/google/generativeai/types/permission_types.py
index 1df831db0..0334ade2a 100644
--- a/google/generativeai/types/permission_types.py
+++ b/google/generativeai/types/permission_types.py
@@ -28,6 +28,7 @@
 from google.generativeai.utils import flatten_update_paths
 from google.generativeai import string_utils
 
+__all__ = ['Permission', 'Permissions']
 
 GranteeType = protos.Permission.GranteeType
 Role = protos.Permission.Role
@@ -89,7 +90,7 @@ def valid_id(name: str) -> bool:
 
 
 @string_utils.prettyprint
-@dataclasses.dataclass
+@dataclasses.dataclass(init=False)
 class Permission:
     """
     A permission to access a resource.
@@ -100,6 +101,24 @@ class Permission:
     grantee_type: Optional[GranteeType]
     email_address: Optional[str] = None
 
+    def __init__(
+        self,
+        name: str,
+        role: RoleOptions,
+        grantee_type: Optional[GranteeTypeOptions] = None,
+        email_address: Optional[str] = None,
+    ):
+        self.name = name
+        if role is None:
+            self.role = None
+        else:
+            self.role = to_role(role)
+        if grantee_type is None:
+            self.grantee_type = None
+        else:
+            self.grantee_type = to_grantee_type(grantee_type)
+        self.email_address = email_address
+
     def delete(
         self,
         client: glm.PermissionServiceClient | None = None,
@@ -133,7 +152,8 @@ def _apply_update(self, path, value):
 
     def update(
         self,
-        updates: dict[str, Any],
+        updates: dict[str, Any]|None,
+        *,
         client: glm.PermissionServiceClient | None = None,
     ) -> Permission:
         """
@@ -279,6 +299,12 @@ def _make_create_permission_request(
                 f"Invalid operation: An 'email_address' must be provided when 'grantee_type' is not set to 'EVERYONE'. Currently, 'grantee_type' is set to '{grantee_type}' and 'email_address' is '{email_address if email_address else 'not provided'}'."
             )
 
+        if email_address and grantee_type is None:
+            if email_address.endswith("googlegroups.com"):
+                grantee_type = GranteeType.GROUP
+            else:
+                grantee_type = GranteeType.USER
+
         permission = protos.Permission(
             role=role,
             grantee_type=grantee_type,
@@ -367,6 +393,10 @@ def list(
             permission = type(permission).to_dict(permission)
             yield Permission(**permission)
 
+    def __iter__(self):
+        return self.list()
+
+
     async def list_async(
         self,
         page_size: Optional[int] = None,
@@ -385,6 +415,42 @@ async def list_async(
             permission = type(permission).to_dict(permission)
             yield Permission(**permission)
 
+    async def __aiter__(self):
+        return await self.async_list()
+
+    @classmethod
+    def get(
+        cls,
+        name: str,
+        client: glm.PermissionServiceClient | None = None,
+    ) -> Permission:
+        """
+        Get information about a specific permission.
+
+        Args:
+            name: The name of the permission to get.
+
+        Returns:
+            Requested permission as an instance of `Permission`.
+        """
+        return Permission.get(name)
+
+    @classmethod
+    async def get_async(
+        cls,
+        name: str
+    ):
+        """
+        Get information about a specific permission.
+
+        Args:
+            name: The name of the permission to get.
+
+        Returns:
+            Requested permission as an instance of `Permission`.
+        """
+        return await Permission.get_async(name)
+
     def transfer_ownership(
         self,
         email_address: str,