From 06eac79c0d5f58c63e0ca23ea1897ad87bd7f585 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=A0ar=C5=ABnas=20Nejus?= <snejus@protonmail.com>
Date: Wed, 4 Sep 2024 04:15:46 +0100
Subject: [PATCH] Centralize requests setup with requests.Session

Improve requests performance with requests.Session which uses connection
pooling for repeated requests to the same host.

Additionally, this centralizes request configuration, making sure that
we use the same timeout and provide beets user agent for all requests.
---
 beetsplug/lyrics.py | 68 +++++++++++++++++----------------------------
 setup.cfg           |  4 +--
 2 files changed, 28 insertions(+), 44 deletions(-)

diff --git a/beetsplug/lyrics.py b/beetsplug/lyrics.py
index 2f8998d96..8e9452a2c 100644
--- a/beetsplug/lyrics.py
+++ b/beetsplug/lyrics.py
@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import atexit
 import errno
 import itertools
 import json
@@ -24,13 +25,12 @@ import os.path
 import re
 import struct
 import unicodedata
-import warnings
 from contextlib import suppress
 from dataclasses import dataclass
 from functools import cached_property, partial, total_ordering
 from http import HTTPStatus
 from typing import TYPE_CHECKING, ClassVar, Iterable, Iterator
-from urllib.parse import quote, urlencode, urlparse
+from urllib.parse import quote, urlparse
 
 import requests
 from typing_extensions import TypedDict
@@ -106,6 +106,22 @@ class NotFoundError(requests.exceptions.HTTPError):
     pass
 
 
+class TimeoutSession(requests.Session):
+    def request(self, *args, **kwargs):
+        kwargs.setdefault("timeout", 10)
+        return super().request(*args, **kwargs)
+
+
+r_session = TimeoutSession()
+r_session.headers.update({"User-Agent": USER_AGENT})
+
+
+@atexit.register
+def close_session():
+    """Close the requests session on shut down."""
+    r_session.close()
+
+
 # Utilities.
 
 
@@ -246,21 +262,7 @@ class Backend:
         is unreachable.
         """
         try:
-            # Disable the InsecureRequestWarning that comes from using
-            # `verify=false`.
-            # https://github.com/kennethreitz/requests/issues/2214
-            # We're not overly worried about the NSA MITMing our lyrics scraper
-            with warnings.catch_warnings():
-                warnings.simplefilter("ignore")
-                r = requests.get(
-                    url,
-                    verify=False,
-                    headers={
-                        "User-Agent": USER_AGENT,
-                    },
-                    timeout=10,
-                    **kwargs,
-                )
+            r = r_session.get(url)
         except requests.RequestException as exc:
             self._log.debug("lyrics request failed: {0}", exc)
             return
@@ -368,9 +370,7 @@ class LRCLib(Backend):
 
     def fetch_json(self, *args, **kwargs):
         """Wrap the request method to raise an exception on HTTP errors."""
-        kwargs.setdefault("timeout", 10)
-        kwargs.setdefault("headers", {"User-Agent": USER_AGENT})
-        r = requests.get(*args, **kwargs)
+        r = r_session.get(*args, **kwargs)
         if r.status_code == HTTPStatus.NOT_FOUND:
             raise NotFoundError("HTTP Error: Not Found", response=r)
         r.raise_for_status()
@@ -535,10 +535,7 @@ class Genius(SearchBackend):
     def __init__(self, config, log):
         super().__init__(config, log)
         self.api_key = config["genius_api_key"].as_str()
-        self.headers = {
-            "Authorization": "Bearer %s" % self.api_key,
-            "User-Agent": USER_AGENT,
-        }
+        self.headers = {"Authorization": f"Bearer {self.api_key}"}
 
     def fetch(self, artist: str, title: str, *_) -> str | None:
         """Fetch lyrics from genius.com
@@ -573,18 +570,13 @@ class Genius(SearchBackend):
         search_url = self.base_url + "/search"
         data = {"q": title + " " + artist.lower()}
         try:
-            response = requests.get(
-                search_url,
-                params=data,
-                headers=self.headers,
-                timeout=10,
-            )
+            r = r_session.get(search_url, params=data, headers=self.headers)
         except requests.RequestException as exc:
             self._log.debug("Genius API request failed: {0}", exc)
             return None
 
         try:
-            return response.json()
+            return r.json()
         except ValueError:
             return None
 
@@ -979,13 +971,7 @@ class LyricsPlugin(plugins.BeetsPlugin):
         }
 
         oauth_url = "https://datamarket.accesscontrol.windows.net/v2/OAuth2-13"
-        oauth_token = json.loads(
-            requests.post(
-                oauth_url,
-                data=urlencode(params),
-                timeout=10,
-            ).content
-        )
+        oauth_token = r_session.post(oauth_url, params=params).json()
         if "access_token" in oauth_token:
             return "Bearer " + oauth_token["access_token"]
         else:
@@ -1202,10 +1188,8 @@ class LyricsPlugin(plugins.BeetsPlugin):
                 "https://api.microsofttranslator.com/v2/Http.svc/"
                 "Translate?text=%s&to=%s" % ("|".join(text_lines), to_lang)
             )
-            r = requests.get(
-                url,
-                headers={"Authorization ": self.bing_auth_token},
-                timeout=10,
+            r = r_session.get(
+                url, headers={"Authorization": self.bing_auth_token}
             )
             if r.status_code != 200:
                 self._log.debug(
diff --git a/setup.cfg b/setup.cfg
index 15ca23f65..8e3d7e3b8 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -21,8 +21,8 @@ omit = beets/test/*
 precision = 2
 skip_empty = true
 show_missing = true
-exclude_lines =
-    pragma: no cover
+exclude_also =
+    @atexit.register
     if TYPE_CHECKING
     if typing.TYPE_CHECKING
     raise AssertionError