From 72ccab8ca04bd1d95d3dd9934928ae813b1c223b Mon Sep 17 00:00:00 2001
From: Stein Magnus Jodal <stein.magnus@jodal.no>
Date: Thu, 16 Nov 2023 02:37:32 +0100
Subject: [PATCH] Extract error handling from client methods

---
 src/brreg/enhetsregisteret/_client.py | 40 +++++++++++++--------------
 1 file changed, 20 insertions(+), 20 deletions(-)

diff --git a/src/brreg/enhetsregisteret/_client.py b/src/brreg/enhetsregisteret/_client.py
index eb95c6b..0d5cede 100644
--- a/src/brreg/enhetsregisteret/_client.py
+++ b/src/brreg/enhetsregisteret/_client.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Optional
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Generator, Optional
 
 import httpx
 
@@ -11,6 +12,22 @@
     from types import TracebackType
 
 
+@contextmanager
+def error_handler() -> Generator[None, Any, None]:
+    try:
+        yield
+    except httpx.HTTPError as exc:
+        response: httpx.Response | None = getattr(exc, "response", None)
+        raise BrregRestError(
+            str(exc),
+            method=(exc.request.method if exc.request else None),
+            url=(str(exc.request.url) if exc.request else None),
+            status_code=(response.status_code if response else None),
+        ) from exc
+    except Exception as exc:
+        raise BrregError(exc) from exc
+
+
 class Client:
     _client: httpx.Client
 
@@ -37,16 +54,8 @@ def close(self) -> None:
         self._client.close()
 
     def get_enhet(self, organisasjonsnummer: str) -> Optional[Enhet]:
-        """Get :class:`Enhet` given an organization number.
-
-        Returns :class:`None` if Enhet is gone or not found.
-        Returns :class:`Enhet` if Enhet is found.
-
-        Raises :class:`BrregRestError` if a REST error occurs.
-        Raises :class:`BrregError` if an unhandled exception occurs.
-        """
-        res: Optional[httpx.Response] = None
-        try:
+        """Get :class:`Enhet` given an organization number."""
+        with error_handler():
             res = self._client.get(
                 f"/enheter/{organisasjonsnummer}",
                 headers={
@@ -57,12 +66,3 @@ def get_enhet(self, organisasjonsnummer: str) -> Optional[Enhet]:
                 return None
             res.raise_for_status()
             return Enhet.model_validate_json(res.content)
-        except httpx.HTTPError as exc:
-            raise BrregRestError(
-                str(exc),
-                method=(exc.request.method if exc.request else None),
-                url=(str(exc.request.url) if exc.request else None),
-                status_code=(res.status_code if res else None),
-            ) from exc
-        except Exception as exc:
-            raise BrregError(exc) from exc