[Downloader] Parse tree URLs when cloning repos (#2119)

* [Downloader] Parse tree URLs when cloning repos

Signed-off-by: Toby <tobyharradine@gmail.com>

* Only match GitHub and GitLab URLs

Signed-off-by: Toby Harradine <tobyharradine@gmail.com>

redbot/cogs/downloader/repo_manager.py

@@ -2,11 +2,12 @@ import asyncio
 import functools
 import os
 import pkgutil
+import re
 from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from subprocess import run as sp_run, PIPE
 from sys import executable
-from typing import Tuple, MutableMapping, Union
+from typing import Tuple, MutableMapping, Union, Optional
 
 from redbot.core import data_manager, commands
 from redbot.core.utils import safe_delete
@@ -489,8 +490,11 @@ class Repo(RepoJSONMixin):
 
 
 class RepoManager:
-    def __init__(self):
 
+    GITHUB_OR_GITLAB_RE = re.compile("https?://git(?:hub)|(?:lab)\.com/")
+    TREE_URL_RE = re.compile(r"(?P<tree>/tree)/(?P<branch>\S+)$")
+
+    def __init__(self):
         self._repos = {}
 
         loop = asyncio.get_event_loop()
@@ -510,7 +514,7 @@ class RepoManager:
             raise InvalidRepoName("Not a valid Python variable name.")
         return name.lower()
 
-    async def add_repo(self, url: str, name: str, branch: str = "master") -> Repo:
+    async def add_repo(self, url: str, name: str, branch: Optional[str] = None) -> Repo:
         """Add and clone a git repository.
 
         Parameters
@@ -533,6 +537,8 @@ class RepoManager:
                 "That repo name you provided already exists. Please choose another."
             )
 
+        url, branch = self._parse_url(url, branch)
+
         # noinspection PyTypeChecker
         r = Repo(url=url, name=name, branch=branch, folder_path=self.repos_folder / name)
         await r.clone()
@@ -630,3 +636,12 @@ class RepoManager:
         if set:
             self._repos = ret
         return ret
+
+    def _parse_url(self, url: str, branch: Optional[str]) -> Tuple[str, Optional[str]]:
+        if self.GITHUB_OR_GITLAB_RE.match(url):
+            tree_url_match = self.TREE_URL_RE.search(url)
+            if tree_url_match:
+                url = url[: tree_url_match.start("tree")]
+                if branch is None:
+                    branch = tree_url_match["branch"]
+            return url, branch

tests/cogs/downloader/test_downloader.py

@@ -94,3 +94,27 @@ async def test_existing_repo(repo_manager):
         await repo_manager.add_repo("http://test.com", "test")
 
     repo_manager.does_repo_exist.assert_called_once_with("test")
+
+
+def test_tree_url_parse(repo_manager):
+    cases = [
+        {
+            "input": ("https://github.com/Tobotimus/Tobo-Cogs", None),
+            "expected": ("https://github.com/Tobotimus/Tobo-Cogs", None),
+        },
+        {
+            "input": ("https://github.com/Tobotimus/Tobo-Cogs", "V3"),
+            "expected": ("https://github.com/Tobotimus/Tobo-Cogs", "V3"),
+        },
+        {
+            "input": ("https://github.com/Tobotimus/Tobo-Cogs/tree/V3", None),
+            "expected": ("https://github.com/Tobotimus/Tobo-Cogs", "V3"),
+        },
+        {
+            "input": ("https://github.com/Tobotimus/Tobo-Cogs/tree/V3", "V4"),
+            "expected": ("https://github.com/Tobotimus/Tobo-Cogs", "V4"),
+        },
+    ]
+
+    for test_case in cases:
+        assert test_case["expected"] == repo_manager._parse_url(*test_case["input"])