From 906f03ff7b8b435733f5837971cd6db83046dddd Mon Sep 17 00:00:00 2001
From: pogzyb <pogzyb@umich.edu>
Date: Mon, 7 Feb 2022 11:10:47 -0500
Subject: [PATCH] Add close and aio_close methods

These methods provide a more direct way to close the underlying httpx_client if the user decides to initialize an RDAPClient.
---
 whodap/__init__.py | 32 ++++++++++++++++----------------
 whodap/client.py   | 14 ++++++++++++++
 2 files changed, 30 insertions(+), 16 deletions(-)

diff --git a/whodap/__init__.py b/whodap/__init__.py
index f407c91..471540f 100644
--- a/whodap/__init__.py
+++ b/whodap/__init__.py
@@ -44,8 +44,8 @@ def lookup_domain(
     finally:
         # close the default client created by DNSClient if it exists;
         # otherwise it's up to the user to close their `httpx_client`
-        if not httpx_client and not dns_client.httpx_client.is_closed:
-            dns_client.httpx_client.close()
+        if not httpx_client:
+            dns_client.close()
     return resp
 
 
@@ -70,8 +70,8 @@ async def aio_lookup_domain(
     finally:
         # close the default client created by DNSClient if it exists;
         # otherwise it's up to the user to close their `httpx_client`
-        if not httpx_client and not dns_client.httpx_client.is_closed:
-            await dns_client.httpx_client.aclose()
+        if not httpx_client:
+            await dns_client.aio_close()
     return resp
 
 
@@ -94,8 +94,8 @@ def lookup_ipv4(
     finally:
         # close the default client created by IPv4Client if it exists;
         # otherwise it's up to the user to close their `httpx_client`
-        if not httpx_client and not ipv4_client.httpx_client.is_closed:
-            ipv4_client.httpx_client.close()
+        if not httpx_client:
+            ipv4_client.close()
     return resp
 
 
@@ -118,8 +118,8 @@ async def aio_lookup_ipv4(
     finally:
         # close the default client created by IPv4Client if it exists;
         # otherwise it's up to the user to close their `httpx_client`
-        if not httpx_client and not ipv4_client.httpx_client.is_closed:
-            await ipv4_client.httpx_client.aclose()
+        if not httpx_client:
+            await ipv4_client.aio_close()
     return resp
 
 
@@ -142,8 +142,8 @@ def lookup_ipv6(
     finally:
         # close the default client created by IPv6Client if it exists;
         # otherwise it's up to the user to close their `httpx_client`
-        if not httpx_client and not ipv6_client.httpx_client.is_closed:
-            ipv6_client.httpx_client.close()
+        if not httpx_client:
+            ipv6_client.close()
     return resp
 
 
@@ -166,8 +166,8 @@ async def aio_lookup_ipv6(
     finally:
         # close the default client created by IPv6Client if it exists;
         # otherwise it's up to the user to close their `httpx_client`
-        if not httpx_client and not ipv6_client.httpx_client.is_closed:
-            await ipv6_client.httpx_client.aclose()
+        if not httpx_client:
+            await ipv6_client.aio_close()
     return resp
 
 
@@ -190,8 +190,8 @@ def lookup_asn(
     finally:
         # close the default client created by ASNClient if it exists;
         # otherwise it's up to the user to close their `httpx_client`
-        if not httpx_client and not asn_client.httpx_client.is_closed:
-            asn_client.httpx_client.close()
+        if not httpx_client:
+            asn_client.close()
     return resp
 
 
@@ -214,6 +214,6 @@ async def aio_lookup_asn(
     finally:
         # close the default client created by ASNClient if it exists;
         # otherwise it's up to the user to close their `httpx_client`
-        if not httpx_client and not asn_client.httpx_client.is_closed:
-            await asn_client.httpx_client.aclose()
+        if not httpx_client:
+            await asn_client.aio_close()
     return resp
diff --git a/whodap/client.py b/whodap/client.py
index d5e53fb..90f5a0f 100644
--- a/whodap/client.py
+++ b/whodap/client.py
@@ -60,6 +60,20 @@ def _set_iana_info(self, iana_resp: Dict[str, Any]) -> None:
         """
         ...
 
+    def close(self):
+        """
+        Closes the underlying `httpx.Client`
+        """
+        if not self.httpx_client.is_closed:
+            self.httpx_client.close()
+
+    async def aio_close(self):
+        """
+        Closes the underlying `httpx.AsyncClient`
+        """
+        if not self.httpx_client.is_closed:
+            await self.httpx_client.aclose()
+
     @classmethod
     @contextmanager
     def new_client_context(cls, httpx_client: Optional[httpx.Client] = None):