Adding support for pyproject.toml file (resolves a part of #32380)

sdks/python/apache_beam/runners/portability/stager.py

@@ -282,10 +282,12 @@ def create_job_resources(
           raise RuntimeError(
               'The file %s cannot be found. It was specified in the '
               '--setup_file command line option.' % setup_options.setup_file)
-        if os.path.basename(setup_options.setup_file) != 'setup.py':
+        if os.path.basename(setup_options.setup_file) not in ('setup.py',
+                                                            'pyproject.toml'):
           raise RuntimeError(
               'The --setup_file option expects the full path to a file named '
-              'setup.py instead of %s' % setup_options.setup_file)
+              'setup.py or pyproject.toml instead of %s' %
+              setup_options.setup_file)
         tarball_file = Stager._build_setup_package(
             setup_options.setup_file, temp_dir, build_setup_args)
         resources.append(
@@ -786,11 +788,13 @@ def _build_setup_package(
       temp_dir: str,
       build_setup_args: Optional[List[str]] = None) -> str:
     saved_current_directory = os.getcwd()
+
     try:
       os.chdir(os.path.dirname(setup_file))
       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.
+        # build the sdist else fallback to legacy
+        # setup.py sdist call for setup.py file.
         try:
           build_setup_args = [
               Stager._get_python_executable(),
@@ -805,15 +809,24 @@ def _build_setup_package(
           _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)
+          if setup_file.endswith('setup.py'):
+            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)
+          else:
+            # If it's pyproject.toml and `python -m build` failed,
+            # there's no direct legacy fallback.
+            raise RuntimeError(
+                f"Failed to build package from '{setup_file}' using . "
+                f"'python -m build'. Please ensure that the 'build' module "
+                f"is installed and your project's build configuration is valid."
+            )
       output_files = glob.glob(os.path.join(temp_dir, '*.tar.gz'))
       if not output_files:
         raise RuntimeError(

sdks/python/apache_beam/runners/portability/stager_test.py

@@ -42,6 +42,10 @@
 
 _LOGGER = logging.getLogger(__name__)
 
+try:
+  import setuptools
+except ImportError:
+  setuptools = None
 
 class StagerTest(unittest.TestCase):
   def setUp(self):
@@ -393,7 +397,46 @@ def test_setup_file_not_named_setup_dot_py(self):
     self.assertTrue(
         cm.exception.args[0].startswith(
             'The --setup_file option expects the full path to a file named '
-            'setup.py instead of '))
+            'setup.py or pyproject.toml instead of '))
+
+  def test_setup_file_not_named_pyproject_dot_toml_file(self):
+    staging_dir = self.make_temp_dir()
+    source_dir = self.make_temp_dir()
+
+    options = PipelineOptions()
+    self.update_options(options)
+    options.view_as(SetupOptions).setup_file = (
+        os.path.join(source_dir, 'xyz-pyproject.toml'))
+
+    self.create_temp_file(
+        os.path.join(source_dir, 'xyz-pyproject.toml'), 'notused')
+    with self.assertRaises(RuntimeError) as cm:
+      self.stager.create_and_stage_job_resources(
+          options, staging_location=staging_dir)
+    self.assertTrue(
+        cm.exception.args[0].startswith(
+            'The --setup_file option expects the full path to a file named '
+            'setup.py or pyproject.toml instead of '))
+
+  @unittest.skipIf(setuptools is None,"setuptools not available in this env")
+  def test_setup_file_as_setup_dot_py_file(self):
+    source_dir = self.make_temp_dir()
+
+    options = PipelineOptions()
+    self.update_options(options)
+    options.view_as(SetupOptions).setup_file = (
+        os.path.join(source_dir, 'setup.py'))
+
+    self.create_temp_file(
+        os.path.join(source_dir, 'setup.py'),
+        "from setuptools import setup; setup(name='my_package')")
+    temp_dir = tempfile.mkdtemp()
+    resources = self.stager.create_job_resources(
+          options=options, temp_dir=temp_dir)
+    file_path = list(self.stager.extract_staging_tuple_iter(resources))[0][0]
+
+    self.assertEqual(os.path.join(temp_dir, 'my_package-0.0.0.tar.gz'),
+                     file_path)
 
   def test_sdk_location_default(self):
     staging_dir = self.make_temp_dir()