From 669e7abb9b12d31d6784bb4df455a71a6481c983 Mon Sep 17 00:00:00 2001 From: moonlightnexus Date: Fri, 5 Apr 2024 23:19:35 +0530 Subject: [PATCH] Added `return_class` parameter to `attach_role`, `remove_role`, and `update_role` methods in `AuthLiteClient` class to allow returning a class instance instead of a dictionary. Added `SignOffSessionReplace` class to represent the response from these methods when `return_class` is True. Raised `ParseError` when `signoff_session_and_assign` is True but `refresh_token` or `access_token` is not provided. --- trustauthx/authlite.py | 79 ++++++++++++++++++++++++++++++++++++++---- trustauthx/scheme.py | 10 ++++++ 2 files changed, 83 insertions(+), 6 deletions(-) diff --git a/trustauthx/authlite.py b/trustauthx/authlite.py index 2796f51..c81387c 100644 --- a/trustauthx/authlite.py +++ b/trustauthx/authlite.py @@ -912,7 +912,26 @@ def _re_init_roles(self) -> _Roles: ) return self.Roles - def attach_role(self, uid:str, rol_ids:str|list, signoff_session_and_assign=False, refresh_token=None, access_token=None): + def attach_role(self, uid:str, rol_ids:str|list, signoff_session_and_assign=False, + refresh_token=None, access_token=None, + return_class:bool=False) -> dict|SignOffSessionReplace: + """ + Attaches a role to a user. + + Args: + uid (str): The user ID to attach the role to. + rol_ids (str | list): The ID(s) of the role(s) to attach. + signoff_session_and_assign (bool, optional): Whether to sign off the session and assign. Default is False. + refresh_token (str, optional): The refresh token for authentication. + access_token (str, optional): The access token for authentication. + return_class (bool, optional): Whether to return a class instance. Default is False. + + Returns: + dict | SignOffSessionReplace: The response from the API, or a class instance if return_class is True. + + Raises: + ParseError: If signoff_session_and_assign is True but refresh_token or access_token is not provided. + """ if signoff_session_and_assign: if not refresh_token or not access_token: raise ParseError( @@ -941,9 +960,31 @@ def attach_role(self, uid:str, rol_ids:str|list, signoff_session_and_assign=Fals "RefreshToken": refresh_token, } response = requests.post(url, headers=headers, params=params, json=data) - return response.json() + if signoff_session_and_assign: return response.json() + else: + if return_class: return SignOffSessionReplace(response.json()) + else: return SignOffSessionReplace(response.json()).to_dict() + + def remove_role(self, uid:str, rol_ids:str|list, signoff_session_and_assign=False, + refresh_token=None, access_token=None, + return_class:bool=False) -> dict|SignOffSessionReplace: + """ + Removes a role from a user. - def remove_role(self, uid:str, rol_ids:str|list, signoff_session_and_assign=False, refresh_token=None, access_token=None): + Args: + uid (str): The user ID to remove the role from. + rol_ids (str | list): The ID(s) of the role(s) to remove. + signoff_session_and_assign (bool, optional): Whether to sign off the session and assign. Default is False. + refresh_token (str, optional): The refresh token for authentication. + access_token (str, optional): The access token for authentication. + return_class (bool, optional): Whether to return a class instance. Default is False. + + Returns: + dict | SignOffSessionReplace: The response from the API, or a class instance if return_class is True. + + Raises: + ParseError: If signoff_session_and_assign is True but refresh_token or access_token is not provided. + """ if signoff_session_and_assign: if not refresh_token or not access_token: raise ParseError( @@ -972,9 +1013,32 @@ def remove_role(self, uid:str, rol_ids:str|list, signoff_session_and_assign=Fals "RefreshToken": refresh_token, } response = requests.post(url, headers=headers, params=params, json=data) - return response.json() + if signoff_session_and_assign: return response.json() + else: + if return_class: return SignOffSessionReplace(response.json()) + else: return SignOffSessionReplace(response.json()).to_dict() + + def update_role(self, uid:str, rol_ids_to_add:str|list, rol_ids_to_remove:str|list, + signoff_session_and_assign=False, refresh_token=None, access_token=None, + return_class:bool=False) -> dict|SignOffSessionReplace: + """ + Updates a user's roles by adding and/or removing roles. + + Args: + uid (str): The user ID to update roles for. + rol_ids_to_add (str | list): The ID(s) of the role(s) to add. + rol_ids_to_remove (str | list): The ID(s) of the role(s) to remove. + signoff_session_and_assign (bool, optional): Whether to sign off the session and assign. Default is False. + refresh_token (str, optional): The refresh token for authentication. + access_token (str, optional): The access token for authentication. + return_class (bool, optional): Whether to return a class instance. Default is False. + + Returns: + dict | SignOffSessionReplace: The response from the API, or a class instance if return_class is True. - def update_role(self, uid:str, rol_ids_to_add:str|list, rol_ids_to_remove:str|list, signoff_session_and_assign=False, refresh_token=None, access_token=None): + Raises: + ParseError: If signoff_session_and_assign is True but refresh_token or access_token is not provided. + """ if signoff_session_and_assign: if not refresh_token or not access_token: raise ParseError( @@ -1007,4 +1071,7 @@ def update_role(self, uid:str, rol_ids_to_add:str|list, rol_ids_to_remove:str|li "RefreshToken": refresh_token, } response = requests.post(url, headers=headers, params=params, json=data) - return response.json() + if signoff_session_and_assign: return response.json() + else: + if return_class: return SignOffSessionReplace(response.json()) + else: return SignOffSessionReplace(response.json()).to_dict() \ No newline at end of file diff --git a/trustauthx/scheme.py b/trustauthx/scheme.py index 547161c..d872c21 100644 --- a/trustauthx/scheme.py +++ b/trustauthx/scheme.py @@ -91,6 +91,16 @@ class User: def to_dict(self): return asdict(self) + +@dataclass +class SignOffSessionReplace: + uid: str + access_token: str + refresh_token: str + role: List[str] + + def to_dict(self): + return asdict(self) """# Demo data demo_get_all_roles_response = GetAllRolesResponse(roles=[