Added an option in setup.py to use any branch from azure-functions-host (#978)

* Added an option in setup.py to use any branch from azure-functions-host

* Addressed comments

Author: gavin-aguiar

setup.py

@@ -25,6 +25,7 @@
 # The GitHub repository of the Azure Functions Host
 WEBHOST_GITHUB_API = "https://api.github.com/repos/Azure/azure-functions-host"
 WEBHOST_TAG_PREFIX = "v4."
+WEBHOST_GIT_REPO = "https://github.com/Azure/azure-functions-host/archive"
 
 # Extensions necessary for non-core bindings.
 AZURE_EXTENSIONS = """\
@@ -56,7 +57,6 @@
 </Project>
 """
 
-
 NUGET_CONFIG = """\
 <?xml version="1.0" encoding="UTF-8"?>
 <configuration>
@@ -75,7 +75,6 @@
 </configuration>
 """
 
-
 CLASSIFIERS = [
     "Development Status :: 5 - Production/Stable",
     'Programming Language :: Python',
@@ -91,7 +90,6 @@
     "Intended Audience :: Developers",
 ]
 
-
 PACKAGES = [
     "azure_functions_worker",
     "azure_functions_worker.protos",
@@ -103,15 +101,13 @@
     "azure_functions_worker._thirdparty"
 ]
 
-
 INSTALL_REQUIRES = [
     "grpcio~=1.43.0",
     "grpcio-tools~=1.43.0",
     "protobuf~=3.19.3",
     "azure-functions==1.9.0"
 ]
 
-
 EXTRA_REQUIRES = {
     "dev": [
         "azure-eventhub~=5.7.0",  # Used for EventHub E2E tests
@@ -135,6 +131,7 @@
 
 class BuildGRPC:
     """Generate gRPC bindings."""
+
     def _gen_grpc(self):
         root = pathlib.Path(os.path.abspath(os.path.dirname(__file__)))
 
@@ -275,16 +272,20 @@ def run(self):
 class Webhost(distutils.cmd.Command):
     description = 'Download and setup Azure Functions Web Host.'
     user_options = [
-        ('webhost-version', None,
-            'A Functions Host version to be downloaded (e.g. 3.0.15278).'),
-        ('webhost-dir', None,
-            'A path to the directory where Azure Web Host will be installed.'),
+        ('webhost-version=', None,
+         'A Functions Host version to be downloaded (e.g. 3.0.15278).'),
+        ('webhost-dir=', None,
+         'A path to the directory where Azure Web Host will be installed.'),
+        ('branch-name=', None,
+         'A branch from where azure-functions-host will be installed '
+         '(e.g. branch-name=dev, branch-name= abc/branchname)')
     ]
 
     def __init__(self, dist: Distribution):
         super().__init__(dist)
         self.webhost_dir = None
         self.webhost_version = None
+        self.branch_name = None
 
     def initialize_options(self):
         pass
@@ -313,14 +314,20 @@ def _get_webhost_version() -> str:
         return latest[0]['name'].replace('v', '')
 
     @staticmethod
-    def _download_webhost_zip(version: str) -> str:
+    def _download_webhost_zip(version: str, branch: str) -> str:
         # Return the path of the downloaded host
         temporary_file = tempfile.NamedTemporaryFile()
-        zip_url = (
-            'https://github.com/Azure/azure-functions-host/archive/'
-            f'v{version}.zip'
-        )
-        print(f'Downloading Functions Host ({version}) from {zip_url}')
+
+        if branch is not None:
+            zip_url = (
+                f'{WEBHOST_GIT_REPO}/refs/heads/{branch}.zip'
+            )
+        else:
+            zip_url = (
+                f'{WEBHOST_GIT_REPO}/v{version}.zip'
+            )
+
+        print(f'Downloading Functions Host from {zip_url}')
 
         with temporary_file as zipf:
             zipf.close()
@@ -410,9 +417,11 @@ def _compile_webhost(webhost_dir: pathlib.Path):
 
     def run(self):
         # Prepare webhost
-        zip_path = self._download_webhost_zip(self.webhost_version)
+        zip_path = self._download_webhost_zip(self.webhost_version,
+                                              self.branch_name)
         self._create_webhost_folder(self.webhost_dir)
-        self._extract_webhost_zip(version=self.webhost_version,
+        version = self.branch_name or self.webhost_version
+        self._extract_webhost_zip(version=version.replace('/', '-'),
                                   src_zip=zip_path,
                                   dest=self.webhost_dir)
         self._chmod_protobuf_generation_script(self.webhost_dir)
@@ -439,18 +448,23 @@ def run(self) -> None:
         for dir_to_delete in self.dir_list_to_delete:
             dir_delete = pathlib.Path(dir_to_delete)
             if dir_delete.exists():
-                dir_util.remove_tree(str(dir_delete))
+                try:
+                    print(f"Deleting directory: {dir_to_delete}")
+                    shutil.rmtree(dir_delete)
+                except OSError as ex:
+                    print(f"Error deleting directory: {dir_to_delete}. "
+                          f"Exception: {ex}")
 
 
 COMMAND_CLASS = {
     'develop': Development,
     'build': BuildProtos,
     'webhost': Webhost,
+    'webhost --branch-name={branch-name}': Webhost,
     'extension': Extension,
     'clean': Clean
 }
 
-
 setup(
     name="azure-functions-worker",
     version=VERSION,