fix: allow import of base sdk without extras installed

See issue: #23

GitOrigin-RevId: 485ef59591db8f79ec82bcda63742a8510447bb7

Author: dmccrystal

assemblyai/extras.py

@@ -1,15 +1,20 @@
 import time
 from typing import Generator
 
-try:
-    import pyaudio
-except ImportError:
-    raise ImportError(
-        "You must install the extras for this SDK to use this feature. "
-        "Run `pip install assemblyai[extras]` to install the extras. "
-        "Make sure to install `apt install portaudio19-dev` (Debian/Ubuntu) or "
-        "`brew install portaudio` (MacOS) before installing the extras."
-    )
+
+class AssemblyAIExtrasNotInstalledError(ImportError):
+    def __init__(
+        self,
+        msg="""
+        You must install the extras for this SDK to use this feature.
+        Run `pip install assemblyai[extras]` to install the extras.
+        Make sure to install `apt install portaudio19-dev` (Debian/Ubuntu) or
+        `brew install portaudio` (MacOS) before installing the extras
+        """,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(msg, *args, **kwargs)
 
 
 class MicrophoneStream:
@@ -25,6 +30,10 @@ def __init__(
             channels: The number of channels to record audio from.
             sample_rate: The sample rate to record audio at.
         """
+        try:
+            import pyaudio
+        except ImportError:
+            raise AssemblyAIExtrasNotInstalledError
 
         self._pyaudio = pyaudio.PyAudio()
         self.sample_rate = sample_rate

tests/unit/test_imports.py

@@ -0,0 +1,61 @@
+import sys
+from importlib import reload
+
+import pytest
+import pytest_mock
+
+import assemblyai as aai
+
+
+class ImportFailureMocker:
+    def __init__(self, module: str):
+        self.module = module
+
+    def find_spec(self, fullname, path, target=None):
+        if fullname == self.module:
+            raise ImportError
+
+    def __enter__(self):
+        # Remove module if already imported
+        if self.module in sys.modules:
+            del sys.modules[self.module]
+
+        # Add self as first importer
+        sys.meta_path.insert(0, self)
+        return self
+
+    def __exit__(self, type, value, traceback):
+        # Remove self as importer
+        sys.meta_path.pop(0)
+
+
+def __reload_assesmblyai_module():
+    reload(aai)
+    aai.settings.api_key = "test"
+
+
+def test_import_sdk_without_extras_installed():
+    with ImportFailureMocker("pyaudio"):
+        __reload_assesmblyai_module()
+        # Test succeeds if no failures
+
+
+def test_import_sdk_and_use_extras_without_extras_installed():
+    with ImportFailureMocker("pyaudio"):
+        __reload_assesmblyai_module()
+
+        with pytest.raises(aai.extras.AssemblyAIExtrasNotInstalledError):
+            aai.extras.MicrophoneStream()
+
+
+def test_import_sdk_and_use_extras_with_extras_installed(
+    mocker: pytest_mock.MockerFixture,
+):
+    import pyaudio
+
+    __reload_assesmblyai_module()
+
+    mocker.patch.object(pyaudio.PyAudio, "open", return_value=None)
+    aai.extras.MicrophoneStream()
+
+    # Test succeeds if no failures

tests/unit/test_realtime_transcriber.py

@@ -13,7 +13,6 @@
 import assemblyai as aai
 
 aai.settings.api_key = "test"
-aai.settings.base_url = "https://api.assemblyai.com/v2"
 
 
 def _disable_rw_threads(mocker: MockFixture):
@@ -28,6 +27,8 @@ def test_realtime_connect_has_parameters(mocker: MockFixture):
     """
     Test that the connect method has the correct parameters set
     """
+    aai.settings.base_url = "https://api.assemblyai.com/v2"
+
     actual_url = None
     actual_additional_headers = None
     actual_open_timeout = None