Skip to content

Commit

Permalink
Add context manager support
Browse files Browse the repository at this point in the history
  • Loading branch information
jshcodes committed Nov 18, 2023
1 parent c258b9f commit 5717ab8
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 0 deletions.
9 changes: 9 additions & 0 deletions src/falconpy/_auth_object/_uber_interface.py
Original file line number Diff line number Diff line change
Expand Up @@ -150,6 +150,15 @@ def logout(self) -> bool:

return bool(result["status_code"] == 200)

def __enter__(self):
"""Allow for entry as a context manager."""
return self

def __exit__(self, *args):
"""Discard our token when we exit the context."""
self.logout()
return args

# Legacy property getters maintained for backwards functionality.
def authenticated(self) -> bool:
"""Return the current authentication status."""
Expand Down
9 changes: 9 additions & 0 deletions src/falconpy/_service_class/_service_class.py
Original file line number Diff line number Diff line change
Expand Up @@ -230,6 +230,15 @@ def override(self,
exp=expand_result
))

def __enter__(self):
"""Allow for entry as a context manager."""
return self

def __exit__(self, *args):
"""Discard our token when we exit the context."""
self.logout()
return args

# ___ ____ ____ ___ ____ ____ ___ _ ____ ____
# |__] |__/ | | |__] |___ |__/ | | |___ [__
# | | \ |__| | |___ | \ | | |___ ___]
Expand Down

0 comments on commit 5717ab8

Please sign in to comment.