use try except instead of module check

apache · Oct 15, 2023 · bb1442e · bb1442e
1 parent 76200ae
commit bb1442e
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/sdks/python/apache_beam/runners/portability/stager.py b/sdks/python/apache_beam/runners/portability/stager.py
@@ -51,7 +51,6 @@
 import hashlib
 import logging
 import os
-import pkgutil
 import shutil
 import sys
 import tempfile
@@ -774,7 +773,7 @@ def _build_setup_package(setup_file,  # type: str
       if build_setup_args is None:
         # if build is installed in the user env, use it to
         # build the sdist else fallback to legacy setup.py sdist call.
-        if pkgutil.find_loader('build') is not None:
+        try:
           build_setup_args = [
               Stager._get_python_executable(),
               '-m',
@@ -784,16 +783,18 @@ def _build_setup_package(setup_file,  # type: str
               temp_dir,
               os.path.dirname(setup_file),
           ]
-        else:
+          _LOGGER.info('Executing command: %s', build_setup_args)
+          processes.check_output(build_setup_args)
+        except RuntimeError:
           build_setup_args = [
               Stager._get_python_executable(),
               os.path.basename(setup_file),
               'sdist',
               '--dist-dir',
               temp_dir
           ]
-      _LOGGER.info('Executing command: %s', build_setup_args)
-      processes.check_output(build_setup_args)
+          _LOGGER.info('Executing command: %s', build_setup_args)
+          processes.check_output(build_setup_args)
       output_files = glob.glob(os.path.join(temp_dir, '*.tar.gz'))
       if not output_files:
         raise RuntimeError(