From 1d011da3e66e4b0a6f8b198d33a086cf2d093894 Mon Sep 17 00:00:00 2001
From: Thomas Newton <thomas.w.newton@gmail.com>
Date: Wed, 7 Feb 2024 20:40:05 +0000
Subject: [PATCH] Python side boilerplate

---
 python/pyarrow/__init__.py | 2 +-
 python/pyarrow/_fs.pyx     | 3 +++
 python/pyarrow/fs.py       | 4 ++++
 python/setup.py            | 6 ++++++
 4 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/python/pyarrow/__init__.py b/python/pyarrow/__init__.py
index 4dbd1258d3cea..cf4f5a3a5abd9 100644
--- a/python/pyarrow/__init__.py
+++ b/python/pyarrow/__init__.py
@@ -145,7 +145,7 @@ def print_entry(label, value):
         print(f"  {module: <20}: {status: <8}")
 
     print("\nFilesystems:")
-    filesystems = ["GcsFileSystem", "HadoopFileSystem", "S3FileSystem"]
+    filesystems = ["AzureFileSystem", "GcsFileSystem", "HadoopFileSystem", "S3FileSystem"]
     for fs in filesystems:
         status = "Enabled" if _filesystem_is_available(fs) else "-"
         print(f"  {fs: <20}: {status: <8}")
diff --git a/python/pyarrow/_fs.pyx b/python/pyarrow/_fs.pyx
index 395f488144331..86cf39e993c1b 100644
--- a/python/pyarrow/_fs.pyx
+++ b/python/pyarrow/_fs.pyx
@@ -491,6 +491,9 @@ cdef class FileSystem(_Weakrefable):
         elif typ == 'gcs':
             from pyarrow._gcsfs import GcsFileSystem
             self = GcsFileSystem.__new__(GcsFileSystem)
+        elif typ == 'abfs':
+            from pyarrow._azurefs import AzureFileSystem
+            self = AzureFileSystem.__new__(AzureFileSystem)
         elif typ == 'hdfs':
             from pyarrow._hdfs import HadoopFileSystem
             self = HadoopFileSystem.__new__(HadoopFileSystem)
diff --git a/python/pyarrow/fs.py b/python/pyarrow/fs.py
index ead750ca44ef8..647419a656f43 100644
--- a/python/pyarrow/fs.py
+++ b/python/pyarrow/fs.py
@@ -39,6 +39,10 @@
 FileStats = FileInfo
 
 _not_imported = []
+try:
+    from pyarrow._azurefs import AzureFileSystem  # noqa
+except ImportError:
+    _not_imported.append("AzureFileSystem")
 
 try:
     from pyarrow._hdfs import HadoopFileSystem  # noqa
diff --git a/python/setup.py b/python/setup.py
index 098d75a3186af..b253d18289a72 100755
--- a/python/setup.py
+++ b/python/setup.py
@@ -105,6 +105,8 @@ def run(self):
                       'build type (debug or release), default release'),
                      ('boost-namespace=', None,
                       'namespace of boost (default: boost)'),
+                     ('with-azure', None, 
+                      'build the Azure Blob Storage extension'),
                      ('with-cuda', None, 'build the Cuda extension'),
                      ('with-flight', None, 'build the Flight extension'),
                      ('with-substrait', None, 'build the Substrait extension'),
@@ -150,6 +152,8 @@ def initialize_options(self):
             if not hasattr(sys, 'gettotalrefcount'):
                 self.build_type = 'release'
 
+        self.with_azure = strtobool(
+            os.environ.get('PYARROW_WITH_AZURE', '0'))
         self.with_gcs = strtobool(
             os.environ.get('PYARROW_WITH_GCS', '0'))
         self.with_s3 = strtobool(
@@ -349,6 +353,8 @@ def _failure_permitted(self, name):
             return True
         if name == '_s3fs' and not self.with_s3:
             return True
+        if name == '_azurefs' and not self.with_azure:
+            return True
         if name == '_hdfs' and not self.with_hdfs:
             return True
         if name == '_dataset' and not self.with_dataset: