Skip to content

Commit 4aa6040

Browse files
Adding support for pyproject.toml file (resolves a part of #32380) (#34924)
* Adding support for pyproject.toml file (resolves a part of #32380) * Corrected Formatting * Update sdks/python/apache_beam/runners/portability/stager_test.py --------- Co-authored-by: tvalentyn <tvalentyn@users.noreply.github.com>
1 parent ad4741c commit 4aa6040

File tree

2 files changed

+70
-13
lines changed

2 files changed

+70
-13
lines changed

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

Lines changed: 25 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -282,10 +282,12 @@ def create_job_resources(
282282
raise RuntimeError(
283283
'The file %s cannot be found. It was specified in the '
284284
'--setup_file command line option.' % setup_options.setup_file)
285-
if os.path.basename(setup_options.setup_file) != 'setup.py':
285+
if os.path.basename(setup_options.setup_file) not in ('setup.py',
286+
'pyproject.toml'):
286287
raise RuntimeError(
287288
'The --setup_file option expects the full path to a file named '
288-
'setup.py instead of %s' % setup_options.setup_file)
289+
'setup.py or pyproject.toml instead of %s' %
290+
setup_options.setup_file)
289291
tarball_file = Stager._build_setup_package(
290292
setup_options.setup_file, temp_dir, build_setup_args)
291293
resources.append(
@@ -786,11 +788,13 @@ def _build_setup_package(
786788
temp_dir: str,
787789
build_setup_args: Optional[List[str]] = None) -> str:
788790
saved_current_directory = os.getcwd()
791+
789792
try:
790793
os.chdir(os.path.dirname(setup_file))
791794
if build_setup_args is None:
792795
# if build is installed in the user env, use it to
793-
# build the sdist else fallback to legacy setup.py sdist call.
796+
# build the sdist else fallback to legacy
797+
# setup.py sdist call for setup.py file.
794798
try:
795799
build_setup_args = [
796800
Stager._get_python_executable(),
@@ -805,15 +809,24 @@ def _build_setup_package(
805809
_LOGGER.info('Executing command: %s', build_setup_args)
806810
processes.check_output(build_setup_args)
807811
except RuntimeError:
808-
build_setup_args = [
809-
Stager._get_python_executable(),
810-
os.path.basename(setup_file),
811-
'sdist',
812-
'--dist-dir',
813-
temp_dir
814-
]
815-
_LOGGER.info('Executing command: %s', build_setup_args)
816-
processes.check_output(build_setup_args)
812+
if setup_file.endswith('setup.py'):
813+
build_setup_args = [
814+
Stager._get_python_executable(),
815+
os.path.basename(setup_file),
816+
'sdist',
817+
'--dist-dir',
818+
temp_dir
819+
]
820+
_LOGGER.info('Executing command: %s', build_setup_args)
821+
processes.check_output(build_setup_args)
822+
else:
823+
# If it's pyproject.toml and `python -m build` failed,
824+
# there's no direct legacy fallback.
825+
raise RuntimeError(
826+
f"Failed to build package from '{setup_file}' using . "
827+
f"'python -m build'. Please ensure that the 'build' module "
828+
f"is installed and your project's build configuration is valid."
829+
)
817830
output_files = glob.glob(os.path.join(temp_dir, '*.tar.gz'))
818831
if not output_files:
819832
raise RuntimeError(

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

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,11 @@
4242

4343
_LOGGER = logging.getLogger(__name__)
4444

45+
try:
46+
import setuptools
47+
except ImportError:
48+
setuptools = None
49+
4550

4651
class StagerTest(unittest.TestCase):
4752
def setUp(self):
@@ -393,7 +398,46 @@ def test_setup_file_not_named_setup_dot_py(self):
393398
self.assertTrue(
394399
cm.exception.args[0].startswith(
395400
'The --setup_file option expects the full path to a file named '
396-
'setup.py instead of '))
401+
'setup.py or pyproject.toml instead of '))
402+
403+
def test_setup_file_supplies_unexpected_filename(self):
404+
staging_dir = self.make_temp_dir()
405+
source_dir = self.make_temp_dir()
406+
407+
options = PipelineOptions()
408+
self.update_options(options)
409+
options.view_as(SetupOptions).setup_file = (
410+
os.path.join(source_dir, 'xyz-pyproject.toml'))
411+
412+
self.create_temp_file(
413+
os.path.join(source_dir, 'xyz-pyproject.toml'), 'notused')
414+
with self.assertRaises(RuntimeError) as cm:
415+
self.stager.create_and_stage_job_resources(
416+
options, staging_location=staging_dir)
417+
self.assertTrue(
418+
cm.exception.args[0].startswith(
419+
'The --setup_file option expects the full path to a file named '
420+
'setup.py or pyproject.toml instead of '))
421+
422+
@unittest.skipIf(setuptools is None, "setuptools not available in this env")
423+
def test_setup_file_as_setup_dot_py_file(self):
424+
source_dir = self.make_temp_dir()
425+
426+
options = PipelineOptions()
427+
self.update_options(options)
428+
options.view_as(SetupOptions).setup_file = (
429+
os.path.join(source_dir, 'setup.py'))
430+
431+
self.create_temp_file(
432+
os.path.join(source_dir, 'setup.py'),
433+
"from setuptools import setup; setup(name='my_package')")
434+
temp_dir = tempfile.mkdtemp()
435+
resources = self.stager.create_job_resources(
436+
options=options, temp_dir=temp_dir)
437+
file_path = list(self.stager.extract_staging_tuple_iter(resources))[0][0]
438+
439+
self.assertEqual(
440+
os.path.join(temp_dir, 'my_package-0.0.0.tar.gz'), file_path)
397441

398442
def test_sdk_location_default(self):
399443
staging_dir = self.make_temp_dir()

0 commit comments

Comments
 (0)