fix: fetch private GitHub URLs with a token on the query string

flake8_nitpick/style.py

@@ -2,6 +2,7 @@
 import logging
 from pathlib import Path
 from typing import TYPE_CHECKING, List, Optional, Set
+from urllib.parse import urlparse, urlunparse
 
 import requests
 import toml
@@ -71,9 +72,6 @@ def get_style_path(self, style_uri: str) -> Optional[Path]:
         """Get the style path from the URI. Add the .toml extension if it's missing."""
         clean_style_uri = style_uri.strip()
 
-        if clean_style_uri and not clean_style_uri.endswith(TOML_EXTENSION):
-            clean_style_uri += TOML_EXTENSION
-
         style_path = None
         if is_url(clean_style_uri) or is_url(self._first_full_path):
             style_path = self.fetch_style_from_url(clean_style_uri)
@@ -86,34 +84,43 @@ def fetch_style_from_url(self, url: str) -> Optional[Path]:
         if self._first_full_path and not is_url(url):
             prefix, rest = self._first_full_path.split(":/")
             resolved = (Path(rest) / url).resolve()
-            url = f"{prefix}:/{resolved}"
+            new_url = f"{prefix}:/{resolved}"
+        else:
+            new_url = url
+
+        parsed_url = list(urlparse(new_url))
+        if not parsed_url[2].endswith(TOML_EXTENSION):
+            parsed_url[2] += TOML_EXTENSION
+        new_url = urlunparse(parsed_url)
 
-        if url in self._already_included:
+        if new_url in self._already_included:
             return None
 
         if not self.config.cache_dir:
             raise FileNotFoundError("Cache dir does not exist")
 
-        response = requests.get(url)
+        response = requests.get(new_url)
         if not response.ok:
-            raise FileNotFoundError(f"Error {response} fetching style URL {url}")
+            raise FileNotFoundError(f"Error {response} fetching style URL {new_url}")
 
         # Save the first full path to be used by the next files without parent.
         if not self._first_full_path:
-            self._first_full_path = url.rsplit("/", 1)[0]
+            self._first_full_path = new_url.rsplit("/", 1)[0]
 
         contents = response.text
-        style_path = self.config.cache_dir / f"{slugify(url)}.toml"
+        style_path = self.config.cache_dir / f"{slugify(new_url)}.toml"
         self.config.cache_dir.mkdir(parents=True, exist_ok=True)
         style_path.write_text(contents)
 
-        LOGGER.info("Loading style from URL %s into %s", url, style_path)
-        self._already_included.add(url)
+        LOGGER.info("Loading style from URL %s into %s", new_url, style_path)
+        self._already_included.add(new_url)
 
         return style_path
 
     def fetch_style_from_local_path(self, partial_file_name: str) -> Optional[Path]:
         """Fetch a style file from a local path."""
+        if partial_file_name and not partial_file_name.endswith(TOML_EXTENSION):
+            partial_file_name += TOML_EXTENSION
         expanded_path = Path(partial_file_name).expanduser()
 
         if not str(expanded_path).startswith("/") and self._first_full_path:

tests/test_style.py

@@ -6,6 +6,7 @@
 
 import responses
 
+from flake8_nitpick.constants import TOML_EXTENSION
 from tests.conftest import TEMP_ROOT_PATH
 from tests.helpers import ProjectMock
 
@@ -361,3 +362,29 @@ def test_relative_style_on_urls(request):
         version = "1.0"
         """
     )
+
+
+@responses.activate
+def test_fetch_private_github_urls(request):
+    """Fetch private GitHub URLs with a token on the query string."""
+    base_url = "https://raw.githubusercontent.com/user/private_repo/branch/path/to/nitpick-style"
+    token = "?token=xxx"
+    full_private_url = f"{base_url}{TOML_EXTENSION}{token}"
+    body = """
+        ["pyproject.toml".tool.black]
+        missing = "thing"
+        """
+    responses.add(responses.GET, f"{full_private_url}", dedent(body), status=200)
+
+    ProjectMock(request).pyproject_toml(
+        f"""
+        [tool.nitpick]
+        style = "{base_url}{token}"
+        """
+    ).lint().assert_single_error(
+        """
+        NIP311 File pyproject.toml has missing values. Use this:
+        [tool.black]
+        missing = "thing"
+    """
+    )