From 3c40b5d251b496874be3d9c989599486a6096db5 Mon Sep 17 00:00:00 2001
From: Juliya Smith <jules@apeworx.io>
Date: Tue, 19 Sep 2023 15:58:40 -0500
Subject: [PATCH] chore: small fixes

---
 ape_safe/_cli.py     | 21 +++++++--------------
 ape_safe/accounts.py | 17 +++++++++++++++++
 2 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/ape_safe/_cli.py b/ape_safe/_cli.py
index bd5d5c9..5c6c662 100644
--- a/ape_safe/_cli.py
+++ b/ape_safe/_cli.py
@@ -11,7 +11,7 @@
 from click import BadArgumentUsage, BadOptionUsage
 
 from ape_safe.accounts import SafeAccount
-from ape_safe.client import ExecutedTxData, SafeClient
+from ape_safe.client import ExecutedTxData
 
 
 @click.group(short_help="Manage Safe accounts and view Safe API data")
@@ -28,14 +28,14 @@ def cli():
 def _list(cli_ctx, network):
     _ = network  # Needed for NetworkBoundCommand
     safes = cli_ctx.account_manager.get_accounts_by_type(type_=SafeAccount)
-    safes_length = len(safes)
+    number_of_safes = len(safes)
 
-    if safes_length == 0:
+    if number_of_safes == 0:
         cli_ctx.logger.warning("No Safes found.")
         return
 
-    header = f"Found {safes_length} Safe"
-    header += "s:" if safes_length > 1 else ":"
+    header = f"Found {number_of_safes} Safe"
+    header += "s:" if number_of_safes > 1 else ":"
     click.echo(header)
 
     for account in safes:
@@ -91,9 +91,7 @@ def remove(cli_ctx, alias):
     safe_container = cli_ctx.account_manager.containers["safe"]
 
     if alias not in safe_container.aliases:
-        raise BadArgumentUsage(
-            f"There is no account with the alias `{alias}` in the safe accounts."
-        )
+        raise BadArgumentUsage(f"There is no safe with the alias `{alias}`.")
 
     address = safe_container.load_account(alias).address
     if click.confirm(f"Remove safe {address} ({alias})"):
@@ -185,12 +183,7 @@ def reject(cli_ctx, network, alias, txn_ids):
 def all_txns(cli_ctx, network, address, confirmed):
     _ = network  # Needed for NetworkBoundCommand
     safe_container = cli_ctx.account_manager.containers["safe"]
-    address = (
-        safe_container.load_account(address).address
-        if address in safe_container.aliases
-        else cli_ctx.conversion_manager.convert(address, AddressType)
-    )
-    client = SafeClient(address=address, chain_id=cli_ctx.chain_manager.provider.chain_id)
+    client = safe_container._get_client(address)
 
     for txn in client.get_transactions(confirmed=confirmed):
         if isinstance(txn, ExecutedTxData):
diff --git a/ape_safe/accounts.py b/ape_safe/accounts.py
index bcdc2ca..bbd8e70 100644
--- a/ape_safe/accounts.py
+++ b/ape_safe/accounts.py
@@ -103,6 +103,23 @@ def delete_account(self, alias: str):
         """
         self._get_path(alias).unlink(missing_ok=True)
 
+    def _get_client(self, key: str) -> BaseSafeClient:
+        if key in self.aliases:
+            safe = self.load_account(key)
+            return safe.client
+
+        elif key in self:
+            account = self[key]  # type: ignore[index]
+            if not isinstance(account, SafeAccount):
+                raise TypeError("Safe container returned non-safe account.")
+
+            return account.client
+
+        else:
+            # Is not locally managed.
+            address = self.conversion_manager.convert(key, AddressType)
+            return SafeClient(address=address, chain_id=self.chain_manager.provider.chain_id)
+
     def _get_path(self, alias: str) -> Path:
         return self.data_folder.joinpath(f"{alias}.json")