From b44549fb30f36f703a327ca2d7addbdc0ff8da5f Mon Sep 17 00:00:00 2001 From: Valentyn Tymofieiev Date: Thu, 29 Jun 2017 10:36:33 -0700 Subject: [PATCH] Simplify package selection logic --- .../apache_beam/runners/dataflow/internal/dependency.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/sdks/python/apache_beam/runners/dataflow/internal/dependency.py b/sdks/python/apache_beam/runners/dataflow/internal/dependency.py index a40a58273d24c..7d36f4f7df95d 100644 --- a/sdks/python/apache_beam/runners/dataflow/internal/dependency.py +++ b/sdks/python/apache_beam/runners/dataflow/internal/dependency.py @@ -545,10 +545,11 @@ def get_sdk_package_name(): """For internal use only; no backwards-compatibility guarantees. Returns the PyPI package name to be staged to Google Cloud Dataflow.""" - container_version = _get_required_container_version() - if container_version == BEAM_CONTAINER_VERSION: + sdk_name, _ = get_sdk_name_and_version() + if sdk_name == 'Google Cloud Dataflow SDK for Python': + return GOOGLE_PACKAGE_NAME + else: return BEAM_PACKAGE_NAME - return GOOGLE_PACKAGE_NAME def _download_pypi_sdk_package(temp_dir):