github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/utils/interactive_utils.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Common interactive utility module."""
    19  # pytype: skip-file
    20  
    21  import logging
    22  
    23  _LOGGER = logging.getLogger(__name__)
    24  
    25  
    26  def is_in_ipython():
    27    """Determines if current code is executed within an ipython session."""
    28    try:
    29      from IPython import get_ipython  # pylint: disable=import-error
    30      if get_ipython():
    31        return True
    32      return False
    33    except ImportError:
    34      # If dependencies are not available, then not interactive for sure.
    35      return False
    36    except (KeyboardInterrupt, SystemExit):
    37      raise
    38    except:  # pylint: disable=bare-except
    39      _LOGGER.info(
    40          'Unexpected error occurred, treated as not in IPython.', exc_info=True)
    41      return False
    42  
    43  
    44  def is_in_notebook():
    45    """Determines if current code is executed from an ipython notebook.
    46  
    47    If is_in_notebook() is True, then is_in_ipython() must also be True.
    48    """
    49    is_in_notebook = False
    50    if is_in_ipython():
    51      # The import and usage must be valid under the execution path.
    52      from IPython import get_ipython
    53      if 'IPKernelApp' in get_ipython().config:
    54        is_in_notebook = True
    55    return is_in_notebook
    56  
    57  
    58  def alter_label_if_ipython(transform, pvalueish):
    59    """Alters the label to an interactive label with ipython prompt metadata
    60    prefixed for the given transform if the given pvalueish belongs to a
    61    user-defined pipeline and current code execution is within an ipython kernel.
    62    Otherwise, noop.
    63  
    64    A label is either a user-defined or auto-generated str name of a PTransform
    65    that is unique within a pipeline. If current environment is_in_ipython(), Beam
    66    can implicitly create interactive labels to replace labels of top-level
    67    PTransforms to be applied. The label is formatted as:
    68    `Cell {prompt}: {original_label}`.
    69    """
    70    if is_in_ipython():
    71      from apache_beam.runners.interactive import interactive_environment as ie
    72      # Tracks user defined pipeline instances in watched scopes so that we only
    73      # alter labels for any transform to pvalueish belonging to those pipeline
    74      # instances, excluding any transform to be applied in other pipeline
    75      # instances the Beam SDK creates implicitly.
    76      ie.current_env().track_user_pipelines()
    77      from IPython import get_ipython
    78      prompt = get_ipython().execution_count
    79      pipeline = _extract_pipeline_of_pvalueish(pvalueish)
    80      if (pipeline
    81          # We only alter for transforms to be applied to user-defined pipelines
    82          # at pipeline construction time.
    83          and pipeline in ie.current_env().tracked_user_pipelines):
    84        transform.label = '[{}]: {}'.format(prompt, transform.label)
    85  
    86  
    87  def _extract_pipeline_of_pvalueish(pvalueish):
    88    """Extracts the pipeline that the given pvalueish belongs to."""
    89    if isinstance(pvalueish, tuple) and len(pvalueish) > 0:
    90      pvalue = pvalueish[0]
    91    elif isinstance(pvalueish, dict) and len(pvalueish) > 0:
    92      pvalue = next(iter(pvalueish.values()))
    93    else:
    94      pvalue = pvalueish
    95    if hasattr(pvalue, 'pipeline'):
    96      return pvalue.pipeline
    97    return None