Google Git
Sign in
apache / beam-site / refs/heads/revert-643-revert-642-updates_release_2.47.0 / . / pydoc / 2.47.0 / _modules / apache_beam / typehints / decorators.html
blob: 5436076be4b46b48c3690d377bdfc858a567a195 [file] [log] [blame]
<!DOCTYPE html>
<!--[if IE 8]><html class="no-js lt-ie9" lang="en" > <![endif]-->
<!--[if gt IE 8]><!--> <html class="no-js" lang="en" > <!--<![endif]-->
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>apache_beam.typehints.decorators &mdash; Apache Beam 2.47.0 documentation</title>
<script type="text/javascript" src="../../../_static/js/modernizr.min.js"></script>
<script type="text/javascript" id="documentation_options" data-url_root="../../../" src="../../../_static/documentation_options.js"></script>
<script type="text/javascript" src="../../../_static/jquery.js"></script>
<script type="text/javascript" src="../../../_static/underscore.js"></script>
<script type="text/javascript" src="../../../_static/doctools.js"></script>
<script type="text/javascript" src="../../../_static/language_data.js"></script>
<script async="async" type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/mathjax/2.7.5/latest.js?config=TeX-AMS-MML_HTMLorMML"></script>
<script type="text/javascript" src="../../../_static/js/theme.js"></script>
<link rel="stylesheet" href="../../../_static/css/theme.css" type="text/css" />
<link rel="stylesheet" href="../../../_static/pygments.css" type="text/css" />
<link rel="index" title="Index" href="../../../genindex.html" />
<link rel="search" title="Search" href="../../../search.html" />
</head>
<body class="wy-body-for-nav">
<div class="wy-grid-for-nav">
<nav data-toggle="wy-nav-shift" class="wy-nav-side">
<div class="wy-side-scroll">
<div class="wy-side-nav-search" >
<a href="../../../index.html" class="icon icon-home"> Apache Beam
</a>
<div class="version">
2.47.0
</div>
<div role="search">
<form id="rtd-search-form" class="wy-form" action="../../../search.html" method="get">
<input type="text" name="q" placeholder="Search docs" />
<input type="hidden" name="check_keywords" value="yes" />
<input type="hidden" name="area" value="default" />
</form>
</div>
</div>
<div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.coders.html">apache_beam.coders package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.dataframe.html">apache_beam.dataframe package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.io.html">apache_beam.io package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.metrics.html">apache_beam.metrics package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.ml.html">apache_beam.ml package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.options.html">apache_beam.options package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.portability.html">apache_beam.portability package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.runners.html">apache_beam.runners package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.testing.html">apache_beam.testing package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.transforms.html">apache_beam.transforms package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.typehints.html">apache_beam.typehints package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.utils.html">apache_beam.utils package</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.yaml.html">apache_beam.yaml package</a></li>
</ul>
<ul>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.error.html">apache_beam.error module</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.pipeline.html">apache_beam.pipeline module</a></li>
<li class="toctree-l1"><a class="reference internal" href="../../../apache_beam.pvalue.html">apache_beam.pvalue module</a></li>
</ul>
</div>
</div>
</nav>
<section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
<nav class="wy-nav-top" aria-label="top navigation">
<i data-toggle="wy-nav-top" class="fa fa-bars"></i>
<a href="../../../index.html">Apache Beam</a>
</nav>
<div class="wy-nav-content">
<div class="rst-content">
<div role="navigation" aria-label="breadcrumbs navigation">
<ul class="wy-breadcrumbs">
<li><a href="../../../index.html">Docs</a> &raquo;</li>
<li><a href="../../index.html">Module code</a> &raquo;</li>
<li>apache_beam.typehints.decorators</li>
<li class="wy-breadcrumbs-aside">
</li>
</ul>
<hr/>
</div>
<div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
<div itemprop="articleBody">
<h1>Source code for apache_beam.typehints.decorators</h1><div class="highlight"><pre>
<span></span><span class="c1">#</span>
<span class="c1"># Licensed to the Apache Software Foundation (ASF) under one or more</span>
<span class="c1"># contributor license agreements. See the NOTICE file distributed with</span>
<span class="c1"># this work for additional information regarding copyright ownership.</span>
<span class="c1"># The ASF licenses this file to You under the Apache License, Version 2.0</span>
<span class="c1"># (the &quot;License&quot;); you may not use this file except in compliance with</span>
<span class="c1"># the License. You may obtain a copy of the License at</span>
<span class="c1">#</span>
<span class="c1"># http://www.apache.org/licenses/LICENSE-2.0</span>
<span class="c1">#</span>
<span class="c1"># Unless required by applicable law or agreed to in writing, software</span>
<span class="c1"># distributed under the License is distributed on an &quot;AS IS&quot; BASIS,</span>
<span class="c1"># WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.</span>
<span class="c1"># See the License for the specific language governing permissions and</span>
<span class="c1"># limitations under the License.</span>
<span class="c1">#</span>
<span class="sd">&quot;&quot;&quot;Type hinting decorators allowing static or runtime type-checking for the SDK.</span>
<span class="sd">This module defines decorators which utilize the type-hints defined in</span>
<span class="sd">&#39;type_hints.py&#39; to allow annotation of the types of function arguments and</span>
<span class="sd">return values.</span>
<span class="sd">Type-hints for functions are annotated using two separate decorators. One is for</span>
<span class="sd">type-hinting the types of function arguments, the other for type-hinting the</span>
<span class="sd">function return value. Type-hints can either be specified in the form of</span>
<span class="sd">positional arguments::</span>
<span class="sd"> @with_input_types(int, int)</span>
<span class="sd"> def add(a, b):</span>
<span class="sd"> return a + b</span>
<span class="sd">Keyword arguments::</span>
<span class="sd"> @with_input_types(a=int, b=int)</span>
<span class="sd"> def add(a, b):</span>
<span class="sd"> return a + b</span>
<span class="sd">Or even a mix of both::</span>
<span class="sd"> @with_input_types(int, b=int)</span>
<span class="sd"> def add(a, b):</span>
<span class="sd"> return a + b</span>
<span class="sd">Example usage for type-hinting arguments only::</span>
<span class="sd"> @with_input_types(s=str)</span>
<span class="sd"> def to_lower(a):</span>
<span class="sd"> return a.lower()</span>
<span class="sd">Example usage for type-hinting return values only::</span>
<span class="sd"> @with_output_types(Tuple[int, bool])</span>
<span class="sd"> def compress_point(ec_point):</span>
<span class="sd"> return ec_point.x, ec_point.y &lt; 0</span>
<span class="sd">Example usage for type-hinting both arguments and return values::</span>
<span class="sd"> @with_input_types(a=int)</span>
<span class="sd"> @with_output_types(str)</span>
<span class="sd"> def int_to_str(a):</span>
<span class="sd"> return str(a)</span>
<span class="sd">The valid type-hint for such as function looks like the following::</span>
<span class="sd"> @with_input_types(a=int, b=int)</span>
<span class="sd"> def foo((a, b)):</span>
<span class="sd"> ...</span>
<span class="sd">Notice that we hint the type of each unpacked argument independently, rather</span>
<span class="sd">than hinting the type of the tuple as a whole (Tuple[int, int]).</span>
<span class="sd">Optionally, type-hints can be type-checked at runtime. To toggle this behavior</span>
<span class="sd">this module defines two functions: &#39;enable_run_time_type_checking&#39; and</span>
<span class="sd">&#39;disable_run_time_type_checking&#39;. NOTE: for this toggle behavior to work</span>
<span class="sd">properly it must appear at the top of the module where all functions are</span>
<span class="sd">defined, or before importing a module containing type-hinted functions.</span>
<span class="sd">&quot;&quot;&quot;</span>
<span class="c1"># pytype: skip-file</span>
<span class="kn">import</span> <span class="nn">inspect</span>
<span class="kn">import</span> <span class="nn">itertools</span>
<span class="kn">import</span> <span class="nn">logging</span>
<span class="kn">import</span> <span class="nn">sys</span>
<span class="kn">import</span> <span class="nn">traceback</span>
<span class="kn">import</span> <span class="nn">types</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Any</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Callable</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Dict</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Iterable</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">NamedTuple</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Optional</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Tuple</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">TypeVar</span>
<span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Union</span>
<span class="kn">from</span> <span class="nn">apache_beam.typehints</span> <span class="kn">import</span> <span class="n">native_type_compatibility</span>
<span class="kn">from</span> <span class="nn">apache_beam.typehints</span> <span class="kn">import</span> <span class="n">typehints</span>
<span class="kn">from</span> <span class="nn">apache_beam.typehints.native_type_compatibility</span> <span class="kn">import</span> <span class="n">convert_to_beam_type</span>
<span class="kn">from</span> <span class="nn">apache_beam.typehints.typehints</span> <span class="kn">import</span> <span class="n">CompositeTypeHintError</span>
<span class="kn">from</span> <span class="nn">apache_beam.typehints.typehints</span> <span class="kn">import</span> <span class="n">SimpleTypeHintError</span>
<span class="kn">from</span> <span class="nn">apache_beam.typehints.typehints</span> <span class="kn">import</span> <span class="n">check_constraint</span>
<span class="kn">from</span> <span class="nn">apache_beam.typehints.typehints</span> <span class="kn">import</span> <span class="n">validate_composite_type_param</span>
<span class="n">__all__</span> <span class="o">=</span> <span class="p">[</span>
<span class="s1">&#39;disable_type_annotations&#39;</span><span class="p">,</span>
<span class="s1">&#39;no_annotations&#39;</span><span class="p">,</span>
<span class="s1">&#39;with_input_types&#39;</span><span class="p">,</span>
<span class="s1">&#39;with_output_types&#39;</span><span class="p">,</span>
<span class="s1">&#39;WithTypeHints&#39;</span><span class="p">,</span>
<span class="s1">&#39;TypeCheckError&#39;</span><span class="p">,</span>
<span class="p">]</span>
<span class="n">T</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">&#39;T&#39;</span><span class="p">)</span>
<span class="n">WithTypeHintsT</span> <span class="o">=</span> <span class="n">TypeVar</span><span class="p">(</span><span class="s1">&#39;WithTypeHintsT&#39;</span><span class="p">,</span> <span class="n">bound</span><span class="o">=</span><span class="s1">&#39;WithTypeHints&#39;</span><span class="p">)</span> <span class="c1"># pylint: disable=invalid-name</span>
<span class="c1"># This is missing in the builtin types module. str.upper is arbitrary, any</span>
<span class="c1"># method on a C-implemented type will do.</span>
<span class="c1"># pylint: disable=invalid-name</span>
<span class="n">_MethodDescriptorType</span> <span class="o">=</span> <span class="nb">type</span><span class="p">(</span><span class="nb">str</span><span class="o">.</span><span class="n">upper</span><span class="p">)</span>
<span class="c1"># pylint: enable=invalid-name</span>
<span class="n">_ANY_VAR_POSITIONAL</span> <span class="o">=</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Tuple</span><span class="p">[</span><span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">,</span> <span class="o">...</span><span class="p">]</span>
<span class="n">_ANY_VAR_KEYWORD</span> <span class="o">=</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Dict</span><span class="p">[</span><span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">,</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">]</span>
<span class="n">_disable_from_callable</span> <span class="o">=</span> <span class="kc">False</span>
<span class="k">def</span> <span class="nf">get_signature</span><span class="p">(</span><span class="n">func</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Like inspect.signature(), but supports Py2 as well.</span>
<span class="sd"> This module uses inspect.signature instead of getfullargspec since in the</span>
<span class="sd"> latter: &#39;the &quot;self&quot; parameter is always reported, even for bound methods&#39;</span>
<span class="sd"> https://github.com/python/cpython/blob/44f91c388a6f4da9ed3300df32ca290b8aa104ea/Lib/inspect.py#L1103</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">signature</span> <span class="o">=</span> <span class="n">inspect</span><span class="o">.</span><span class="n">signature</span><span class="p">(</span><span class="n">func</span><span class="p">)</span>
<span class="k">except</span> <span class="ne">ValueError</span><span class="p">:</span>
<span class="c1"># Fall back on a catch-all signature.</span>
<span class="n">params</span> <span class="o">=</span> <span class="p">[</span>
<span class="n">inspect</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="s1">&#39;_&#39;</span><span class="p">,</span> <span class="n">inspect</span><span class="o">.</span><span class="n">Parameter</span><span class="o">.</span><span class="n">POSITIONAL_OR_KEYWORD</span><span class="p">),</span>
<span class="n">inspect</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span>
<span class="s1">&#39;__unknown__varargs&#39;</span><span class="p">,</span> <span class="n">inspect</span><span class="o">.</span><span class="n">Parameter</span><span class="o">.</span><span class="n">VAR_POSITIONAL</span><span class="p">),</span>
<span class="n">inspect</span><span class="o">.</span><span class="n">Parameter</span><span class="p">(</span><span class="s1">&#39;__unknown__keywords&#39;</span><span class="p">,</span> <span class="n">inspect</span><span class="o">.</span><span class="n">Parameter</span><span class="o">.</span><span class="n">VAR_KEYWORD</span><span class="p">)</span>
<span class="p">]</span>
<span class="n">signature</span> <span class="o">=</span> <span class="n">inspect</span><span class="o">.</span><span class="n">Signature</span><span class="p">(</span><span class="n">params</span><span class="p">)</span>
<span class="c1"># This is a specialization to hint the first argument of certain builtins,</span>
<span class="c1"># such as str.strip.</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">func</span><span class="p">,</span> <span class="n">_MethodDescriptorType</span><span class="p">):</span>
<span class="n">params</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">signature</span><span class="o">.</span><span class="n">parameters</span><span class="o">.</span><span class="n">values</span><span class="p">())</span>
<span class="k">if</span> <span class="n">params</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">annotation</span> <span class="o">==</span> <span class="n">params</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">empty</span><span class="p">:</span>
<span class="n">params</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">params</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">annotation</span><span class="o">=</span><span class="n">func</span><span class="o">.</span><span class="vm">__objclass__</span><span class="p">)</span>
<span class="n">signature</span> <span class="o">=</span> <span class="n">signature</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">parameters</span><span class="o">=</span><span class="n">params</span><span class="p">)</span>
<span class="c1"># This is a specialization to hint the return value of type callables.</span>
<span class="k">if</span> <span class="p">(</span><span class="n">signature</span><span class="o">.</span><span class="n">return_annotation</span> <span class="o">==</span> <span class="n">signature</span><span class="o">.</span><span class="n">empty</span> <span class="ow">and</span>
<span class="nb">isinstance</span><span class="p">(</span><span class="n">func</span><span class="p">,</span> <span class="nb">type</span><span class="p">)):</span>
<span class="n">signature</span> <span class="o">=</span> <span class="n">signature</span><span class="o">.</span><span class="n">replace</span><span class="p">(</span><span class="n">return_annotation</span><span class="o">=</span><span class="n">typehints</span><span class="o">.</span><span class="n">normalize</span><span class="p">(</span><span class="n">func</span><span class="p">))</span>
<span class="k">return</span> <span class="n">signature</span>
<div class="viewcode-block" id="no_annotations"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.no_annotations">[docs]</a><span class="k">def</span> <span class="nf">no_annotations</span><span class="p">(</span><span class="n">fn</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Decorator that prevents Beam from using type hint annotations on a</span>
<span class="sd"> callable.&quot;&quot;&quot;</span>
<span class="nb">setattr</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="s1">&#39;_beam_no_annotations&#39;</span><span class="p">,</span> <span class="kc">True</span><span class="p">)</span>
<span class="k">return</span> <span class="n">fn</span></div>
<div class="viewcode-block" id="disable_type_annotations"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.disable_type_annotations">[docs]</a><span class="k">def</span> <span class="nf">disable_type_annotations</span><span class="p">():</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Prevent Beam from using type hint annotations to determine input and output</span>
<span class="sd"> types of transforms.</span>
<span class="sd"> This setting applies globally.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">global</span> <span class="n">_disable_from_callable</span>
<span class="n">_disable_from_callable</span> <span class="o">=</span> <span class="kc">True</span></div>
<span class="n">TRACEBACK_LIMIT</span> <span class="o">=</span> <span class="mi">5</span>
<span class="k">class</span> <span class="nc">IOTypeHints</span><span class="p">(</span><span class="n">NamedTuple</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Encapsulates all type hint information about a Beam construct.</span>
<span class="sd"> This should primarily be used via the WithTypeHints mixin class, though</span>
<span class="sd"> may also be attached to other objects (such as Python functions).</span>
<span class="sd"> Attributes:</span>
<span class="sd"> input_types: (tuple, dict) List of typing types, and an optional dictionary.</span>
<span class="sd"> May be None. The list and dict correspond to args and kwargs.</span>
<span class="sd"> output_types: (tuple, dict) List of typing types, and an optional dictionary</span>
<span class="sd"> (unused). Only the first element of the list is used. May be None.</span>
<span class="sd"> origin: (List[str]) Stack of tracebacks of method calls used to create this</span>
<span class="sd"> instance.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="n">input_types</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">Any</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]]</span>
<span class="n">output_types</span><span class="p">:</span> <span class="n">Optional</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">Tuple</span><span class="p">[</span><span class="n">Any</span><span class="p">,</span> <span class="o">...</span><span class="p">],</span> <span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">Any</span><span class="p">]]]</span>
<span class="n">origin</span><span class="p">:</span> <span class="n">List</span><span class="p">[</span><span class="nb">str</span><span class="p">]</span>
<span class="nd">@classmethod</span>
<span class="k">def</span> <span class="nf">_make_origin</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">bases</span><span class="p">,</span> <span class="n">tb</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="n">msg</span><span class="o">=</span><span class="p">()):</span>
<span class="c1"># type: (List[IOTypeHints], bool, Iterable[str]) -&gt; List[str]</span>
<span class="k">if</span> <span class="n">msg</span><span class="p">:</span>
<span class="n">res</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">res</span> <span class="o">=</span> <span class="p">[]</span>
<span class="k">if</span> <span class="n">tb</span><span class="p">:</span>
<span class="c1"># Omit this method and the IOTypeHints method that called it.</span>
<span class="n">num_frames_skip</span> <span class="o">=</span> <span class="mi">2</span>
<span class="n">tbs</span> <span class="o">=</span> <span class="n">traceback</span><span class="o">.</span><span class="n">format_stack</span><span class="p">(</span><span class="n">limit</span><span class="o">=</span><span class="n">TRACEBACK_LIMIT</span> <span class="o">+</span>
<span class="n">num_frames_skip</span><span class="p">)[:</span><span class="o">-</span><span class="n">num_frames_skip</span><span class="p">]</span>
<span class="c1"># tb is a list of strings in the form of &#39;File ...\n[code]\n&#39;. Split into</span>
<span class="c1"># single lines and flatten.</span>
<span class="n">res</span> <span class="o">+=</span> <span class="nb">list</span><span class="p">(</span>
<span class="n">itertools</span><span class="o">.</span><span class="n">chain</span><span class="o">.</span><span class="n">from_iterable</span><span class="p">(</span><span class="n">s</span><span class="o">.</span><span class="n">strip</span><span class="p">()</span><span class="o">.</span><span class="n">split</span><span class="p">(</span><span class="s1">&#39;</span><span class="se">\n</span><span class="s1">&#39;</span><span class="p">)</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">tbs</span><span class="p">))</span>
<span class="n">bases</span> <span class="o">=</span> <span class="p">[</span><span class="n">base</span> <span class="k">for</span> <span class="n">base</span> <span class="ow">in</span> <span class="n">bases</span> <span class="k">if</span> <span class="n">base</span><span class="o">.</span><span class="n">origin</span><span class="p">]</span>
<span class="k">if</span> <span class="n">bases</span><span class="p">:</span>
<span class="n">res</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39;&#39;</span><span class="p">,</span> <span class="s1">&#39;based on:&#39;</span><span class="p">]</span>
<span class="k">for</span> <span class="n">i</span><span class="p">,</span> <span class="n">base</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">bases</span><span class="p">):</span>
<span class="k">if</span> <span class="n">i</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
<span class="n">res</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39;&#39;</span><span class="p">,</span> <span class="s1">&#39;and:&#39;</span><span class="p">]</span>
<span class="n">res</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39; &#39;</span> <span class="o">+</span> <span class="nb">str</span><span class="p">(</span><span class="n">base</span><span class="p">)]</span>
<span class="n">res</span> <span class="o">+=</span> <span class="p">[</span><span class="s1">&#39; &#39;</span> <span class="o">+</span> <span class="n">s</span> <span class="k">for</span> <span class="n">s</span> <span class="ow">in</span> <span class="n">base</span><span class="o">.</span><span class="n">origin</span><span class="p">]</span>
<span class="k">return</span> <span class="n">res</span>
<span class="nd">@classmethod</span>
<span class="k">def</span> <span class="nf">empty</span><span class="p">(</span><span class="bp">cls</span><span class="p">):</span>
<span class="c1"># type: () -&gt; IOTypeHints</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Construct a base IOTypeHints object with no hints.&quot;&quot;&quot;</span>
<span class="k">return</span> <span class="n">IOTypeHints</span><span class="p">(</span><span class="kc">None</span><span class="p">,</span> <span class="kc">None</span><span class="p">,</span> <span class="p">[])</span>
<span class="nd">@classmethod</span>
<span class="k">def</span> <span class="nf">from_callable</span><span class="p">(</span><span class="bp">cls</span><span class="p">,</span> <span class="n">fn</span><span class="p">):</span>
<span class="c1"># type: (Callable) -&gt; Optional[IOTypeHints]</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Construct an IOTypeHints object from a callable&#39;s signature.</span>
<span class="sd"> Supports Python 3 annotations. For partial annotations, sets unknown types</span>
<span class="sd"> to Any, _ANY_VAR_POSITIONAL, or _ANY_VAR_KEYWORD.</span>
<span class="sd"> Returns:</span>
<span class="sd"> A new IOTypeHints or None if no annotations found.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="n">_disable_from_callable</span> <span class="ow">or</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="s1">&#39;_beam_no_annotations&#39;</span><span class="p">,</span> <span class="kc">False</span><span class="p">):</span>
<span class="k">return</span> <span class="kc">None</span>
<span class="n">signature</span> <span class="o">=</span> <span class="n">get_signature</span><span class="p">(</span><span class="n">fn</span><span class="p">)</span>
<span class="k">if</span> <span class="p">(</span><span class="nb">all</span><span class="p">(</span><span class="n">param</span><span class="o">.</span><span class="n">annotation</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">empty</span>
<span class="k">for</span> <span class="n">param</span> <span class="ow">in</span> <span class="n">signature</span><span class="o">.</span><span class="n">parameters</span><span class="o">.</span><span class="n">values</span><span class="p">())</span> <span class="ow">and</span>
<span class="n">signature</span><span class="o">.</span><span class="n">return_annotation</span> <span class="o">==</span> <span class="n">signature</span><span class="o">.</span><span class="n">empty</span><span class="p">):</span>
<span class="k">return</span> <span class="kc">None</span>
<span class="n">input_args</span> <span class="o">=</span> <span class="p">[]</span>
<span class="n">input_kwargs</span> <span class="o">=</span> <span class="p">{}</span>
<span class="k">for</span> <span class="n">param</span> <span class="ow">in</span> <span class="n">signature</span><span class="o">.</span><span class="n">parameters</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
<span class="k">if</span> <span class="n">param</span><span class="o">.</span><span class="n">annotation</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">empty</span><span class="p">:</span>
<span class="k">if</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">VAR_POSITIONAL</span><span class="p">:</span>
<span class="n">input_args</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">_ANY_VAR_POSITIONAL</span><span class="p">)</span>
<span class="k">elif</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">VAR_KEYWORD</span><span class="p">:</span>
<span class="n">input_kwargs</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">_ANY_VAR_KEYWORD</span>
<span class="k">elif</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">KEYWORD_ONLY</span><span class="p">:</span>
<span class="n">input_kwargs</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Any</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">input_args</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">if</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="ow">in</span> <span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">KEYWORD_ONLY</span><span class="p">,</span> <span class="n">param</span><span class="o">.</span><span class="n">VAR_KEYWORD</span><span class="p">]:</span>
<span class="n">input_kwargs</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">convert_to_beam_type</span><span class="p">(</span><span class="n">param</span><span class="o">.</span><span class="n">annotation</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">assert</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="ow">in</span> <span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">POSITIONAL_ONLY</span><span class="p">,</span>
<span class="n">param</span><span class="o">.</span><span class="n">POSITIONAL_OR_KEYWORD</span><span class="p">,</span>
<span class="n">param</span><span class="o">.</span><span class="n">VAR_POSITIONAL</span><span class="p">],</span> \
<span class="s1">&#39;Unsupported Parameter kind: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span>
<span class="n">input_args</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">convert_to_beam_type</span><span class="p">(</span><span class="n">param</span><span class="o">.</span><span class="n">annotation</span><span class="p">))</span>
<span class="n">output_args</span> <span class="o">=</span> <span class="p">[]</span>
<span class="k">if</span> <span class="n">signature</span><span class="o">.</span><span class="n">return_annotation</span> <span class="o">!=</span> <span class="n">signature</span><span class="o">.</span><span class="n">empty</span><span class="p">:</span>
<span class="n">output_args</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">convert_to_beam_type</span><span class="p">(</span><span class="n">signature</span><span class="o">.</span><span class="n">return_annotation</span><span class="p">))</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">output_args</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">)</span>
<span class="n">name</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="s1">&#39;__name__&#39;</span><span class="p">,</span> <span class="s1">&#39;&lt;unknown&gt;&#39;</span><span class="p">)</span>
<span class="n">msg</span> <span class="o">=</span> <span class="p">[</span><span class="s1">&#39;from_callable(</span><span class="si">%s</span><span class="s1">)&#39;</span> <span class="o">%</span> <span class="n">name</span><span class="p">,</span> <span class="s1">&#39; signature: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">signature</span><span class="p">]</span>
<span class="k">if</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="s1">&#39;__code__&#39;</span><span class="p">):</span>
<span class="n">msg</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
<span class="s1">&#39; File &quot;</span><span class="si">%s</span><span class="s1">&quot;, line </span><span class="si">%d</span><span class="s1">&#39;</span> <span class="o">%</span>
<span class="p">(</span><span class="n">fn</span><span class="o">.</span><span class="vm">__code__</span><span class="o">.</span><span class="n">co_filename</span><span class="p">,</span> <span class="n">fn</span><span class="o">.</span><span class="vm">__code__</span><span class="o">.</span><span class="n">co_firstlineno</span><span class="p">))</span>
<span class="k">return</span> <span class="n">IOTypeHints</span><span class="p">(</span>
<span class="n">input_types</span><span class="o">=</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">input_args</span><span class="p">),</span> <span class="n">input_kwargs</span><span class="p">),</span>
<span class="n">output_types</span><span class="o">=</span><span class="p">(</span><span class="nb">tuple</span><span class="p">(</span><span class="n">output_args</span><span class="p">),</span> <span class="p">{}),</span>
<span class="n">origin</span><span class="o">=</span><span class="bp">cls</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([],</span> <span class="n">tb</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">msg</span><span class="o">=</span><span class="n">msg</span><span class="p">))</span>
<span class="k">def</span> <span class="nf">with_input_types</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
<span class="c1"># type: (...) -&gt; IOTypeHints</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_replace</span><span class="p">(</span>
<span class="n">input_types</span><span class="o">=</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">),</span> <span class="n">origin</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([</span><span class="bp">self</span><span class="p">]))</span>
<span class="k">def</span> <span class="nf">with_output_types</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">args</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
<span class="c1"># type: (...) -&gt; IOTypeHints</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_replace</span><span class="p">(</span>
<span class="n">output_types</span><span class="o">=</span><span class="p">(</span><span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span><span class="p">),</span> <span class="n">origin</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([</span><span class="bp">self</span><span class="p">]))</span>
<span class="k">def</span> <span class="nf">with_input_types_from</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">other</span><span class="p">):</span>
<span class="c1"># type: (IOTypeHints) -&gt; IOTypeHints</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_replace</span><span class="p">(</span>
<span class="n">input_types</span><span class="o">=</span><span class="n">other</span><span class="o">.</span><span class="n">input_types</span><span class="p">,</span> <span class="n">origin</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([</span><span class="bp">self</span><span class="p">]))</span>
<span class="k">def</span> <span class="nf">with_output_types_from</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">other</span><span class="p">):</span>
<span class="c1"># type: (IOTypeHints) -&gt; IOTypeHints</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_replace</span><span class="p">(</span>
<span class="n">output_types</span><span class="o">=</span><span class="n">other</span><span class="o">.</span><span class="n">output_types</span><span class="p">,</span> <span class="n">origin</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([</span><span class="bp">self</span><span class="p">]))</span>
<span class="k">def</span> <span class="nf">simple_output_type</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">context</span><span class="p">):</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_has_output_types</span><span class="p">():</span>
<span class="n">args</span><span class="p">,</span> <span class="n">kwargs</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">args</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">1</span> <span class="ow">or</span> <span class="n">kwargs</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span>
<span class="s1">&#39;Expected single output type hint for </span><span class="si">%s</span><span class="s1"> but got: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span>
<span class="p">(</span><span class="n">context</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">))</span>
<span class="k">return</span> <span class="n">args</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">def</span> <span class="nf">has_simple_output_type</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Whether there&#39;s a single positional output type.&quot;&quot;&quot;</span>
<span class="k">return</span> <span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">output_types</span> <span class="ow">and</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span>
<span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">[</span><span class="mi">1</span><span class="p">])</span>
<span class="k">def</span> <span class="nf">strip_pcoll</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="kn">from</span> <span class="nn">apache_beam.pipeline</span> <span class="kn">import</span> <span class="n">Pipeline</span>
<span class="kn">from</span> <span class="nn">apache_beam.pvalue</span> <span class="kn">import</span> <span class="n">PBegin</span>
<span class="kn">from</span> <span class="nn">apache_beam.pvalue</span> <span class="kn">import</span> <span class="n">PDone</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">strip_pcoll_helper</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_types</span><span class="p">,</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_has_input_types</span><span class="p">,</span>
<span class="s1">&#39;input_types&#39;</span><span class="p">,</span>
<span class="p">[</span><span class="n">Pipeline</span><span class="p">,</span> <span class="n">PBegin</span><span class="p">],</span>
<span class="s1">&#39;This input type hint will be ignored &#39;</span>
<span class="s1">&#39;and not used for type-checking purposes. &#39;</span>
<span class="s1">&#39;Typically, input type hints for a &#39;</span>
<span class="s1">&#39;PTransform are single (or nested) types &#39;</span>
<span class="s1">&#39;wrapped by a PCollection, or PBegin.&#39;</span><span class="p">,</span>
<span class="s1">&#39;strip_pcoll_input()&#39;</span><span class="p">)</span><span class="o">.</span>\
<span class="n">strip_pcoll_helper</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">,</span>
<span class="bp">self</span><span class="o">.</span><span class="n">has_simple_output_type</span><span class="p">,</span>
<span class="s1">&#39;output_types&#39;</span><span class="p">,</span>
<span class="p">[</span><span class="n">PDone</span><span class="p">,</span> <span class="kc">None</span><span class="p">],</span>
<span class="s1">&#39;This output type hint will be ignored &#39;</span>
<span class="s1">&#39;and not used for type-checking purposes. &#39;</span>
<span class="s1">&#39;Typically, output type hints for a &#39;</span>
<span class="s1">&#39;PTransform are single (or nested) types &#39;</span>
<span class="s1">&#39;wrapped by a PCollection, PDone, or None.&#39;</span><span class="p">,</span>
<span class="s1">&#39;strip_pcoll_output()&#39;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">strip_pcoll_helper</span><span class="p">(</span>
<span class="bp">self</span><span class="p">,</span>
<span class="n">my_type</span><span class="p">,</span> <span class="c1"># type: any</span>
<span class="n">has_my_type</span><span class="p">,</span> <span class="c1"># type: Callable[[], bool]</span>
<span class="n">my_key</span><span class="p">,</span> <span class="c1"># type: str</span>
<span class="n">special_containers</span><span class="p">,</span> <span class="c1"># type: List[Union[PBegin, PDone, PCollection]] # noqa: F821</span>
<span class="n">error_str</span><span class="p">,</span> <span class="c1"># type: str</span>
<span class="n">source_str</span> <span class="c1"># type: str</span>
<span class="p">):</span>
<span class="c1"># type: (...) -&gt; IOTypeHints</span>
<span class="kn">from</span> <span class="nn">apache_beam.pvalue</span> <span class="kn">import</span> <span class="n">PCollection</span>
<span class="k">if</span> <span class="ow">not</span> <span class="n">has_my_type</span><span class="p">()</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">my_type</span> <span class="ow">or</span> <span class="nb">len</span><span class="p">(</span><span class="n">my_type</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>
<span class="k">return</span> <span class="bp">self</span>
<span class="n">my_type</span> <span class="o">=</span> <span class="n">my_type</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">my_type</span><span class="p">,</span> <span class="n">typehints</span><span class="o">.</span><span class="n">AnyTypeConstraint</span><span class="p">):</span>
<span class="k">return</span> <span class="bp">self</span>
<span class="n">special_containers</span> <span class="o">+=</span> <span class="p">[</span><span class="n">PCollection</span><span class="p">]</span>
<span class="n">kwarg_dict</span> <span class="o">=</span> <span class="p">{}</span>
<span class="k">if</span> <span class="p">(</span><span class="n">my_type</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">special_containers</span> <span class="ow">and</span>
<span class="nb">getattr</span><span class="p">(</span><span class="n">my_type</span><span class="p">,</span> <span class="s1">&#39;__origin__&#39;</span><span class="p">,</span> <span class="kc">None</span><span class="p">)</span> <span class="o">!=</span> <span class="n">PCollection</span><span class="p">):</span>
<span class="n">logging</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="n">error_str</span> <span class="o">+</span> <span class="s1">&#39; Got: </span><span class="si">%s</span><span class="s1"> instead.&#39;</span> <span class="o">%</span> <span class="n">my_type</span><span class="p">)</span>
<span class="n">kwarg_dict</span><span class="p">[</span><span class="n">my_key</span><span class="p">]</span> <span class="o">=</span> <span class="kc">None</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_replace</span><span class="p">(</span>
<span class="n">origin</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([</span><span class="bp">self</span><span class="p">],</span> <span class="n">tb</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">msg</span><span class="o">=</span><span class="p">[</span><span class="n">source_str</span><span class="p">]),</span>
<span class="o">**</span><span class="n">kwarg_dict</span><span class="p">)</span>
<span class="k">if</span> <span class="p">(</span><span class="nb">getattr</span><span class="p">(</span><span class="n">my_type</span><span class="p">,</span> <span class="s1">&#39;__args__&#39;</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">)</span> <span class="ow">in</span> <span class="p">[</span><span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="kc">None</span><span class="p">]</span> <span class="ow">or</span>
<span class="nb">len</span><span class="p">(</span><span class="n">my_type</span><span class="o">.</span><span class="n">__args__</span><span class="p">)</span> <span class="o">==</span> <span class="mi">0</span><span class="p">):</span>
<span class="c1"># e.g. PCollection (or PBegin/PDone)</span>
<span class="n">kwarg_dict</span><span class="p">[</span><span class="n">my_key</span><span class="p">]</span> <span class="o">=</span> <span class="p">((</span><span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">,</span> <span class="p">),</span> <span class="p">{})</span>
<span class="k">else</span><span class="p">:</span>
<span class="c1"># e.g. PCollection[type]</span>
<span class="n">kwarg_dict</span><span class="p">[</span><span class="n">my_key</span><span class="p">]</span> <span class="o">=</span> <span class="p">((</span><span class="n">convert_to_beam_type</span><span class="p">(</span><span class="n">my_type</span><span class="o">.</span><span class="n">__args__</span><span class="p">[</span><span class="mi">0</span><span class="p">]),</span> <span class="p">),</span> <span class="p">{})</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_replace</span><span class="p">(</span>
<span class="n">origin</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([</span><span class="bp">self</span><span class="p">],</span> <span class="n">tb</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">msg</span><span class="o">=</span><span class="p">[</span><span class="n">source_str</span><span class="p">]),</span>
<span class="o">**</span><span class="n">kwarg_dict</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">strip_iterable</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="c1"># type: () -&gt; IOTypeHints</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Removes outer Iterable (or equivalent) from output type.</span>
<span class="sd"> Only affects instances with simple output types, otherwise is a no-op.</span>
<span class="sd"> Does not modify self.</span>
<span class="sd"> Designed to be used with type hints from callables of ParDo, FlatMap, DoFn.</span>
<span class="sd"> Output type may be Optional[T], in which case the result of stripping T is</span>
<span class="sd"> used as the output type.</span>
<span class="sd"> Output type may be None/NoneType, in which case nothing is done.</span>
<span class="sd"> Example: Generator[Tuple(int, int)] becomes Tuple(int, int)</span>
<span class="sd"> Returns:</span>
<span class="sd"> A copy of this instance with a possibly different output type.</span>
<span class="sd"> Raises:</span>
<span class="sd"> ValueError if output type is simple and not iterable.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="ow">not</span> <span class="bp">self</span><span class="o">.</span><span class="n">has_simple_output_type</span><span class="p">():</span>
<span class="k">return</span> <span class="bp">self</span>
<span class="n">output_type</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">[</span><span class="mi">0</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>
<span class="k">if</span> <span class="n">output_type</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">output_type</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="kc">None</span><span class="p">)):</span>
<span class="k">return</span> <span class="bp">self</span>
<span class="c1"># If output_type == Optional[T]: output_type = T.</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">output_type</span><span class="p">,</span> <span class="n">typehints</span><span class="o">.</span><span class="n">UnionConstraint</span><span class="p">):</span>
<span class="n">types</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">output_type</span><span class="o">.</span><span class="n">union_types</span><span class="p">)</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">types</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span><span class="p">:</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">types</span><span class="o">.</span><span class="n">remove</span><span class="p">(</span><span class="nb">type</span><span class="p">(</span><span class="kc">None</span><span class="p">))</span>
<span class="n">output_type</span> <span class="o">=</span> <span class="n">types</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">except</span> <span class="ne">ValueError</span><span class="p">:</span>
<span class="k">pass</span>
<span class="n">yielded_type</span> <span class="o">=</span> <span class="n">typehints</span><span class="o">.</span><span class="n">get_yielded_type</span><span class="p">(</span><span class="n">output_type</span><span class="p">)</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_replace</span><span class="p">(</span>
<span class="n">output_types</span><span class="o">=</span><span class="p">((</span><span class="n">yielded_type</span><span class="p">,</span> <span class="p">),</span> <span class="p">{}),</span>
<span class="n">origin</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([</span><span class="bp">self</span><span class="p">],</span> <span class="n">tb</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">msg</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;strip_iterable()&#39;</span><span class="p">]))</span>
<span class="k">def</span> <span class="nf">with_defaults</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">hints</span><span class="p">):</span>
<span class="c1"># type: (Optional[IOTypeHints]) -&gt; IOTypeHints</span>
<span class="k">if</span> <span class="ow">not</span> <span class="n">hints</span><span class="p">:</span>
<span class="k">return</span> <span class="bp">self</span>
<span class="k">if</span> <span class="ow">not</span> <span class="bp">self</span><span class="p">:</span>
<span class="k">return</span> <span class="n">hints</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_has_input_types</span><span class="p">():</span>
<span class="n">input_types</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_types</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">input_types</span> <span class="o">=</span> <span class="n">hints</span><span class="o">.</span><span class="n">input_types</span>
<span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">_has_output_types</span><span class="p">():</span>
<span class="n">output_types</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">output_types</span> <span class="o">=</span> <span class="n">hints</span><span class="o">.</span><span class="n">output_types</span>
<span class="n">res</span> <span class="o">=</span> <span class="n">IOTypeHints</span><span class="p">(</span>
<span class="n">input_types</span><span class="p">,</span>
<span class="n">output_types</span><span class="p">,</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_make_origin</span><span class="p">([</span><span class="bp">self</span><span class="p">,</span> <span class="n">hints</span><span class="p">],</span> <span class="n">tb</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span> <span class="n">msg</span><span class="o">=</span><span class="p">[</span><span class="s1">&#39;with_defaults()&#39;</span><span class="p">]))</span>
<span class="k">if</span> <span class="n">res</span> <span class="o">==</span> <span class="bp">self</span><span class="p">:</span>
<span class="k">return</span> <span class="bp">self</span> <span class="c1"># Don&#39;t needlessly increase origin traceback length.</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">return</span> <span class="n">res</span>
<span class="k">def</span> <span class="nf">_has_input_types</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_types</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">any</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_types</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">_has_output_types</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="ow">and</span> <span class="nb">any</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">)</span>
<span class="k">def</span> <span class="fm">__bool__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_has_input_types</span><span class="p">()</span> <span class="ow">or</span> <span class="bp">self</span><span class="o">.</span><span class="n">_has_output_types</span><span class="p">()</span>
<span class="k">def</span> <span class="fm">__repr__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">return</span> <span class="s1">&#39;IOTypeHints[inputs=</span><span class="si">%s</span><span class="s1">, outputs=</span><span class="si">%s</span><span class="s1">]&#39;</span> <span class="o">%</span> <span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">input_types</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">debug_str</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">return</span> <span class="s1">&#39;</span><span class="se">\n</span><span class="s1">&#39;</span><span class="o">.</span><span class="n">join</span><span class="p">([</span><span class="bp">self</span><span class="o">.</span><span class="fm">__repr__</span><span class="p">()]</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">origin</span><span class="p">)</span>
<span class="k">def</span> <span class="fm">__eq__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">other</span><span class="p">):</span>
<span class="k">def</span> <span class="nf">same</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">b</span><span class="p">):</span>
<span class="k">if</span> <span class="n">a</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">a</span><span class="p">):</span>
<span class="k">return</span> <span class="n">b</span> <span class="ow">is</span> <span class="kc">None</span> <span class="ow">or</span> <span class="ow">not</span> <span class="nb">any</span><span class="p">(</span><span class="n">b</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="k">return</span> <span class="n">a</span> <span class="o">==</span> <span class="n">b</span>
<span class="k">return</span> <span class="p">(</span>
<span class="n">same</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_types</span><span class="p">,</span> <span class="n">other</span><span class="o">.</span><span class="n">input_types</span><span class="p">)</span> <span class="ow">and</span>
<span class="n">same</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">,</span> <span class="n">other</span><span class="o">.</span><span class="n">output_types</span><span class="p">))</span>
<span class="k">def</span> <span class="fm">__hash__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">return</span> <span class="nb">hash</span><span class="p">(</span><span class="nb">str</span><span class="p">(</span><span class="bp">self</span><span class="p">))</span>
<span class="k">def</span> <span class="nf">__reduce__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="c1"># Don&#39;t include &quot;origin&quot; debug information in pickled form.</span>
<span class="k">return</span> <span class="p">(</span><span class="n">IOTypeHints</span><span class="p">,</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">input_types</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">output_types</span><span class="p">,</span> <span class="p">[]))</span>
<div class="viewcode-block" id="WithTypeHints"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.WithTypeHints">[docs]</a><span class="k">class</span> <span class="nc">WithTypeHints</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;A mixin class that provides the ability to set and retrieve type hints.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">unused_args</span><span class="p">,</span> <span class="o">**</span><span class="n">unused_kwargs</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_type_hints</span> <span class="o">=</span> <span class="n">IOTypeHints</span><span class="o">.</span><span class="n">empty</span><span class="p">()</span>
<span class="k">def</span> <span class="nf">_get_or_create_type_hints</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="c1"># type: () -&gt; IOTypeHints</span>
<span class="c1"># __init__ may have not been called</span>
<span class="k">try</span><span class="p">:</span>
<span class="c1"># Only return an instance bound to self (see BEAM-8629).</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="vm">__dict__</span><span class="p">[</span><span class="s1">&#39;_type_hints&#39;</span><span class="p">]</span>
<span class="k">except</span> <span class="ne">KeyError</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_type_hints</span> <span class="o">=</span> <span class="n">IOTypeHints</span><span class="o">.</span><span class="n">empty</span><span class="p">()</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">_type_hints</span>
<div class="viewcode-block" id="WithTypeHints.get_type_hints"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.WithTypeHints.get_type_hints">[docs]</a> <span class="k">def</span> <span class="nf">get_type_hints</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Gets and/or initializes type hints for this object.</span>
<span class="sd"> If type hints have not been set, attempts to initialize type hints in this</span>
<span class="sd"> order:</span>
<span class="sd"> - Using self.default_type_hints().</span>
<span class="sd"> - Using self.__class__ type hints.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">return</span> <span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_get_or_create_type_hints</span><span class="p">()</span><span class="o">.</span><span class="n">with_defaults</span><span class="p">(</span>
<span class="bp">self</span><span class="o">.</span><span class="n">default_type_hints</span><span class="p">())</span><span class="o">.</span><span class="n">with_defaults</span><span class="p">(</span>
<span class="n">get_type_hints</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="vm">__class__</span><span class="p">)))</span></div>
<span class="k">def</span> <span class="nf">_set_type_hints</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">type_hints</span><span class="p">):</span>
<span class="c1"># type: (IOTypeHints) -&gt; None</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_type_hints</span> <span class="o">=</span> <span class="n">type_hints</span>
<div class="viewcode-block" id="WithTypeHints.default_type_hints"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.WithTypeHints.default_type_hints">[docs]</a> <span class="k">def</span> <span class="nf">default_type_hints</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">return</span> <span class="kc">None</span></div>
<div class="viewcode-block" id="WithTypeHints.with_input_types"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.WithTypeHints.with_input_types">[docs]</a> <span class="k">def</span> <span class="nf">with_input_types</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">arg_hints</span><span class="p">,</span> <span class="o">**</span><span class="n">kwarg_hints</span><span class="p">):</span>
<span class="c1"># type: (WithTypeHintsT, *Any, **Any) -&gt; WithTypeHintsT</span>
<span class="n">arg_hints</span> <span class="o">=</span> <span class="n">native_type_compatibility</span><span class="o">.</span><span class="n">convert_to_beam_types</span><span class="p">(</span><span class="n">arg_hints</span><span class="p">)</span>
<span class="n">kwarg_hints</span> <span class="o">=</span> <span class="n">native_type_compatibility</span><span class="o">.</span><span class="n">convert_to_beam_types</span><span class="p">(</span><span class="n">kwarg_hints</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_type_hints</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_or_create_type_hints</span><span class="p">()</span><span class="o">.</span><span class="n">with_input_types</span><span class="p">(</span>
<span class="o">*</span><span class="n">arg_hints</span><span class="p">,</span> <span class="o">**</span><span class="n">kwarg_hints</span><span class="p">)</span>
<span class="k">return</span> <span class="bp">self</span></div>
<div class="viewcode-block" id="WithTypeHints.with_output_types"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.WithTypeHints.with_output_types">[docs]</a> <span class="k">def</span> <span class="nf">with_output_types</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">*</span><span class="n">arg_hints</span><span class="p">,</span> <span class="o">**</span><span class="n">kwarg_hints</span><span class="p">):</span>
<span class="c1"># type: (WithTypeHintsT, *Any, **Any) -&gt; WithTypeHintsT</span>
<span class="n">arg_hints</span> <span class="o">=</span> <span class="n">native_type_compatibility</span><span class="o">.</span><span class="n">convert_to_beam_types</span><span class="p">(</span><span class="n">arg_hints</span><span class="p">)</span>
<span class="n">kwarg_hints</span> <span class="o">=</span> <span class="n">native_type_compatibility</span><span class="o">.</span><span class="n">convert_to_beam_types</span><span class="p">(</span><span class="n">kwarg_hints</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">_type_hints</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_get_or_create_type_hints</span><span class="p">()</span><span class="o">.</span><span class="n">with_output_types</span><span class="p">(</span>
<span class="o">*</span><span class="n">arg_hints</span><span class="p">,</span> <span class="o">**</span><span class="n">kwarg_hints</span><span class="p">)</span>
<span class="k">return</span> <span class="bp">self</span></div></div>
<div class="viewcode-block" id="TypeCheckError"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.TypeCheckError">[docs]</a><span class="k">class</span> <span class="nc">TypeCheckError</span><span class="p">(</span><span class="ne">Exception</span><span class="p">):</span>
<span class="k">pass</span></div>
<span class="k">def</span> <span class="nf">_positional_arg_hints</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="n">hints</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Returns the type of a (possibly tuple-packed) positional argument.</span>
<span class="sd"> E.g. for lambda ((a, b), c): None the single positional argument is (as</span>
<span class="sd"> returned by inspect) [[a, b], c] which should have type</span>
<span class="sd"> Tuple[Tuple[Int, Any], float] when applied to the type hints</span>
<span class="sd"> {a: int, b: Any, c: float}.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
<span class="k">return</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Tuple</span><span class="p">[[</span><span class="n">_positional_arg_hints</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">hints</span><span class="p">)</span> <span class="k">for</span> <span class="n">a</span> <span class="ow">in</span> <span class="n">arg</span><span class="p">]]</span>
<span class="k">return</span> <span class="n">hints</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">_unpack_positional_arg_hints</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="n">hint</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Unpacks the given hint according to the nested structure of arg.</span>
<span class="sd"> For example, if arg is [[a, b], c] and hint is Tuple[Any, int], then</span>
<span class="sd"> this function would return ((Any, Any), int) so it can be used in conjunction</span>
<span class="sd"> with inspect.getcallargs.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="nb">list</span><span class="p">):</span>
<span class="n">tuple_constraint</span> <span class="o">=</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Tuple</span><span class="p">[[</span><span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">]</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">arg</span><span class="p">)]</span>
<span class="k">if</span> <span class="ow">not</span> <span class="n">typehints</span><span class="o">.</span><span class="n">is_consistent_with</span><span class="p">(</span><span class="n">hint</span><span class="p">,</span> <span class="n">tuple_constraint</span><span class="p">):</span>
<span class="k">raise</span> <span class="n">TypeCheckError</span><span class="p">(</span>
<span class="s1">&#39;Bad tuple arguments for </span><span class="si">%s</span><span class="s1">: expected </span><span class="si">%s</span><span class="s1">, got </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span>
<span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="n">tuple_constraint</span><span class="p">,</span> <span class="n">hint</span><span class="p">))</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">hint</span><span class="p">,</span> <span class="n">typehints</span><span class="o">.</span><span class="n">TupleConstraint</span><span class="p">):</span>
<span class="k">return</span> <span class="nb">tuple</span><span class="p">(</span>
<span class="n">_unpack_positional_arg_hints</span><span class="p">(</span><span class="n">a</span><span class="p">,</span> <span class="n">t</span><span class="p">)</span> <span class="k">for</span> <span class="n">a</span><span class="p">,</span>
<span class="n">t</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">arg</span><span class="p">,</span> <span class="n">hint</span><span class="o">.</span><span class="n">tuple_types</span><span class="p">))</span>
<span class="k">return</span> <span class="p">(</span><span class="n">typehints</span><span class="o">.</span><span class="n">Any</span><span class="p">,</span> <span class="p">)</span> <span class="o">*</span> <span class="nb">len</span><span class="p">(</span><span class="n">arg</span><span class="p">)</span>
<span class="k">return</span> <span class="n">hint</span>
<span class="k">def</span> <span class="nf">_normalize_var_positional_hint</span><span class="p">(</span><span class="n">hint</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Converts a var_positional hint into Tuple[Union[&lt;types&gt;], ...] form.</span>
<span class="sd"> Args:</span>
<span class="sd"> hint: (tuple) Should be either a tuple of one or more types, or a single</span>
<span class="sd"> Tuple[&lt;type&gt;, ...].</span>
<span class="sd"> Raises:</span>
<span class="sd"> TypeCheckError if hint does not have the right form.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="ow">not</span> <span class="n">hint</span> <span class="ow">or</span> <span class="nb">type</span><span class="p">(</span><span class="n">hint</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">tuple</span><span class="p">:</span>
<span class="k">raise</span> <span class="n">TypeCheckError</span><span class="p">(</span><span class="s1">&#39;Unexpected VAR_POSITIONAL value: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">hint</span><span class="p">)</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">hint</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">hint</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">typehints</span><span class="o">.</span><span class="n">TupleSequenceConstraint</span><span class="p">):</span>
<span class="c1"># Example: tuple(Tuple[Any, ...]) -&gt; Tuple[Any, ...]</span>
<span class="k">return</span> <span class="n">hint</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">else</span><span class="p">:</span>
<span class="c1"># Example: tuple(int, str) -&gt; Tuple[Union[int, str], ...]</span>
<span class="k">return</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Tuple</span><span class="p">[</span><span class="n">typehints</span><span class="o">.</span><span class="n">Union</span><span class="p">[</span><span class="n">hint</span><span class="p">],</span> <span class="o">...</span><span class="p">]</span>
<span class="k">def</span> <span class="nf">_normalize_var_keyword_hint</span><span class="p">(</span><span class="n">hint</span><span class="p">,</span> <span class="n">arg_name</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Converts a var_keyword hint into Dict[&lt;key type&gt;, &lt;value type&gt;] form.</span>
<span class="sd"> Args:</span>
<span class="sd"> hint: (dict) Should either contain a pair (arg_name,</span>
<span class="sd"> Dict[&lt;key type&gt;, &lt;value type&gt;]), or one or more possible types for the</span>
<span class="sd"> value.</span>
<span class="sd"> arg_name: (str) The keyword receiving this hint.</span>
<span class="sd"> Raises:</span>
<span class="sd"> TypeCheckError if hint does not have the right form.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="ow">not</span> <span class="n">hint</span> <span class="ow">or</span> <span class="nb">type</span><span class="p">(</span><span class="n">hint</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">dict</span><span class="p">:</span>
<span class="k">raise</span> <span class="n">TypeCheckError</span><span class="p">(</span><span class="s1">&#39;Unexpected VAR_KEYWORD value: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">hint</span><span class="p">)</span>
<span class="n">keys</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">hint</span><span class="o">.</span><span class="n">keys</span><span class="p">())</span>
<span class="n">values</span> <span class="o">=</span> <span class="nb">list</span><span class="p">(</span><span class="n">hint</span><span class="o">.</span><span class="n">values</span><span class="p">())</span>
<span class="k">if</span> <span class="p">(</span><span class="nb">len</span><span class="p">(</span><span class="n">values</span><span class="p">)</span> <span class="o">==</span> <span class="mi">1</span> <span class="ow">and</span> <span class="n">keys</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="o">==</span> <span class="n">arg_name</span> <span class="ow">and</span>
<span class="nb">isinstance</span><span class="p">(</span><span class="n">values</span><span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="n">typehints</span><span class="o">.</span><span class="n">DictConstraint</span><span class="p">)):</span>
<span class="c1"># Example: dict(kwargs=Dict[str, Any]) -&gt; Dict[str, Any]</span>
<span class="k">return</span> <span class="n">values</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
<span class="k">else</span><span class="p">:</span>
<span class="c1"># Example: dict(k1=str, k2=int) -&gt; Dict[str, Union[str,int]]</span>
<span class="k">return</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Union</span><span class="p">[</span><span class="n">values</span><span class="p">]]</span>
<span class="k">def</span> <span class="nf">getcallargs_forhints</span><span class="p">(</span><span class="n">func</span><span class="p">,</span> <span class="o">*</span><span class="n">type_args</span><span class="p">,</span> <span class="o">**</span><span class="n">type_kwargs</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Bind type_args and type_kwargs to func.</span>
<span class="sd"> Works like inspect.getcallargs, with some modifications to support type hint</span>
<span class="sd"> checks.</span>
<span class="sd"> For unbound args, will use annotations and fall back to Any (or variants of</span>
<span class="sd"> Any).</span>
<span class="sd"> Returns:</span>
<span class="sd"> A mapping from parameter name to argument.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">signature</span> <span class="o">=</span> <span class="n">get_signature</span><span class="p">(</span><span class="n">func</span><span class="p">)</span>
<span class="k">except</span> <span class="ne">ValueError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
<span class="n">logging</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="s1">&#39;Could not get signature for function: </span><span class="si">%s</span><span class="s1">: </span><span class="si">%s</span><span class="s1">&#39;</span><span class="p">,</span> <span class="n">func</span><span class="p">,</span> <span class="n">e</span><span class="p">)</span>
<span class="k">return</span> <span class="p">{}</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">bindings</span> <span class="o">=</span> <span class="n">signature</span><span class="o">.</span><span class="n">bind</span><span class="p">(</span><span class="o">*</span><span class="n">type_args</span><span class="p">,</span> <span class="o">**</span><span class="n">type_kwargs</span><span class="p">)</span>
<span class="k">except</span> <span class="ne">TypeError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
<span class="c1"># Might be raised due to too few or too many arguments.</span>
<span class="k">raise</span> <span class="n">TypeCheckError</span><span class="p">(</span><span class="n">e</span><span class="p">)</span>
<span class="n">bound_args</span> <span class="o">=</span> <span class="n">bindings</span><span class="o">.</span><span class="n">arguments</span>
<span class="k">for</span> <span class="n">param</span> <span class="ow">in</span> <span class="n">signature</span><span class="o">.</span><span class="n">parameters</span><span class="o">.</span><span class="n">values</span><span class="p">():</span>
<span class="k">if</span> <span class="n">param</span><span class="o">.</span><span class="n">name</span> <span class="ow">in</span> <span class="n">bound_args</span><span class="p">:</span>
<span class="c1"># Bound: unpack/convert variadic arguments.</span>
<span class="k">if</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">VAR_POSITIONAL</span><span class="p">:</span>
<span class="n">bound_args</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">_normalize_var_positional_hint</span><span class="p">(</span>
<span class="n">bound_args</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">])</span>
<span class="k">elif</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">VAR_KEYWORD</span><span class="p">:</span>
<span class="n">bound_args</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">_normalize_var_keyword_hint</span><span class="p">(</span>
<span class="n">bound_args</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">],</span> <span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
<span class="k">else</span><span class="p">:</span>
<span class="c1"># Unbound: must have a default or be variadic.</span>
<span class="k">if</span> <span class="n">param</span><span class="o">.</span><span class="n">annotation</span> <span class="o">!=</span> <span class="n">param</span><span class="o">.</span><span class="n">empty</span><span class="p">:</span>
<span class="n">bound_args</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">param</span><span class="o">.</span><span class="n">annotation</span>
<span class="k">elif</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">VAR_POSITIONAL</span><span class="p">:</span>
<span class="n">bound_args</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">_ANY_VAR_POSITIONAL</span>
<span class="k">elif</span> <span class="n">param</span><span class="o">.</span><span class="n">kind</span> <span class="o">==</span> <span class="n">param</span><span class="o">.</span><span class="n">VAR_KEYWORD</span><span class="p">:</span>
<span class="n">bound_args</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">_ANY_VAR_KEYWORD</span>
<span class="k">elif</span> <span class="n">param</span><span class="o">.</span><span class="n">default</span> <span class="ow">is</span> <span class="ow">not</span> <span class="n">param</span><span class="o">.</span><span class="n">empty</span><span class="p">:</span>
<span class="c1"># Declare unbound parameters with defaults to be Any.</span>
<span class="n">bound_args</span><span class="p">[</span><span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">]</span> <span class="o">=</span> <span class="n">typehints</span><span class="o">.</span><span class="n">Any</span>
<span class="k">else</span><span class="p">:</span>
<span class="c1"># This case should be caught by signature.bind() above.</span>
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s1">&#39;Unexpected unbound parameter: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="n">param</span><span class="o">.</span><span class="n">name</span><span class="p">)</span>
<span class="k">return</span> <span class="nb">dict</span><span class="p">(</span><span class="n">bound_args</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">get_type_hints</span><span class="p">(</span><span class="n">fn</span><span class="p">):</span>
<span class="c1"># type: (Any) -&gt; IOTypeHints</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Gets the type hint associated with an arbitrary object fn.</span>
<span class="sd"> Always returns a valid IOTypeHints object, creating one if necessary.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="c1"># pylint: disable=protected-access</span>
<span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="s1">&#39;_type_hints&#39;</span><span class="p">):</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">fn</span><span class="o">.</span><span class="n">_type_hints</span> <span class="o">=</span> <span class="n">IOTypeHints</span><span class="o">.</span><span class="n">empty</span><span class="p">()</span>
<span class="k">except</span> <span class="p">(</span><span class="ne">AttributeError</span><span class="p">,</span> <span class="ne">TypeError</span><span class="p">):</span>
<span class="c1"># Can&#39;t add arbitrary attributes to this object,</span>
<span class="c1"># but might have some restrictions anyways...</span>
<span class="n">hints</span> <span class="o">=</span> <span class="n">IOTypeHints</span><span class="o">.</span><span class="n">empty</span><span class="p">()</span>
<span class="c1"># Python 3.7 introduces annotations for _MethodDescriptorTypes.</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">fn</span><span class="p">,</span> <span class="n">_MethodDescriptorType</span><span class="p">)</span> <span class="ow">and</span> <span class="n">sys</span><span class="o">.</span><span class="n">version_info</span> <span class="o">&lt;</span> <span class="p">(</span><span class="mi">3</span><span class="p">,</span> <span class="mi">7</span><span class="p">):</span>
<span class="n">hints</span> <span class="o">=</span> <span class="n">hints</span><span class="o">.</span><span class="n">with_input_types</span><span class="p">(</span><span class="n">fn</span><span class="o">.</span><span class="vm">__objclass__</span><span class="p">)</span> <span class="c1"># type: ignore</span>
<span class="k">return</span> <span class="n">hints</span>
<span class="k">return</span> <span class="n">fn</span><span class="o">.</span><span class="n">_type_hints</span>
<span class="c1"># pylint: enable=protected-access</span>
<div class="viewcode-block" id="with_input_types"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.with_input_types">[docs]</a><span class="k">def</span> <span class="nf">with_input_types</span><span class="p">(</span><span class="o">*</span><span class="n">positional_hints</span><span class="p">,</span> <span class="o">**</span><span class="n">keyword_hints</span><span class="p">):</span>
<span class="c1"># type: (*Any, **Any) -&gt; Callable[[T], T]</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;A decorator that type-checks defined type-hints with passed func arguments.</span>
<span class="sd"> All type-hinted arguments can be specified using positional arguments,</span>
<span class="sd"> keyword arguments, or a mix of both. Additionaly, all function arguments must</span>
<span class="sd"> be type-hinted in totality if even one parameter is type-hinted.</span>
<span class="sd"> Once fully decorated, if the arguments passed to the resulting function</span>
<span class="sd"> violate the type-hint constraints defined, a :class:`TypeCheckError`</span>
<span class="sd"> detailing the error will be raised.</span>
<span class="sd"> To be used as:</span>
<span class="sd"> .. testcode::</span>
<span class="sd"> from apache_beam.typehints import with_input_types</span>
<span class="sd"> @with_input_types(str)</span>
<span class="sd"> def upper(s):</span>
<span class="sd"> return s.upper()</span>
<span class="sd"> Or:</span>
<span class="sd"> .. testcode::</span>
<span class="sd"> from apache_beam.typehints import with_input_types</span>
<span class="sd"> from apache_beam.typehints import List</span>
<span class="sd"> from apache_beam.typehints import Tuple</span>
<span class="sd"> @with_input_types(ls=List[Tuple[int, int]])</span>
<span class="sd"> def increment(ls):</span>
<span class="sd"> [(i + 1, j + 1) for (i,j) in ls]</span>
<span class="sd"> Args:</span>
<span class="sd"> *positional_hints: Positional type-hints having identical order as the</span>
<span class="sd"> function&#39;s formal arguments. Values for this argument must either be a</span>
<span class="sd"> built-in Python type or an instance of a</span>
<span class="sd"> :class:`~apache_beam.typehints.typehints.TypeConstraint` created by</span>
<span class="sd"> &#39;indexing&#39; a</span>
<span class="sd"> :class:`~apache_beam.typehints.typehints.CompositeTypeHint` instance</span>
<span class="sd"> with a type parameter.</span>
<span class="sd"> **keyword_hints: Keyword arguments mirroring the names of the parameters to</span>
<span class="sd"> the decorated functions. The value of each keyword argument must either</span>
<span class="sd"> be one of the allowed built-in Python types, a custom class, or an</span>
<span class="sd"> instance of a :class:`~apache_beam.typehints.typehints.TypeConstraint`</span>
<span class="sd"> created by &#39;indexing&#39; a</span>
<span class="sd"> :class:`~apache_beam.typehints.typehints.CompositeTypeHint` instance</span>
<span class="sd"> with a type parameter.</span>
<span class="sd"> Raises:</span>
<span class="sd"> :class:`ValueError`: If not all function arguments have</span>
<span class="sd"> corresponding type-hints specified. Or if the inner wrapper function isn&#39;t</span>
<span class="sd"> passed a function object.</span>
<span class="sd"> :class:`TypeCheckError`: If the any of the passed type-hint</span>
<span class="sd"> constraints are not a type or</span>
<span class="sd"> :class:`~apache_beam.typehints.typehints.TypeConstraint` instance.</span>
<span class="sd"> Returns:</span>
<span class="sd"> The original function decorated such that it enforces type-hint constraints</span>
<span class="sd"> for all received function arguments.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="n">converted_positional_hints</span> <span class="o">=</span> <span class="p">(</span>
<span class="n">native_type_compatibility</span><span class="o">.</span><span class="n">convert_to_beam_types</span><span class="p">(</span><span class="n">positional_hints</span><span class="p">))</span>
<span class="n">converted_keyword_hints</span> <span class="o">=</span> <span class="p">(</span>
<span class="n">native_type_compatibility</span><span class="o">.</span><span class="n">convert_to_beam_types</span><span class="p">(</span><span class="n">keyword_hints</span><span class="p">))</span>
<span class="k">del</span> <span class="n">positional_hints</span>
<span class="k">del</span> <span class="n">keyword_hints</span>
<span class="k">def</span> <span class="nf">annotate_input_types</span><span class="p">(</span><span class="n">f</span><span class="p">):</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="n">types</span><span class="o">.</span><span class="n">FunctionType</span><span class="p">):</span>
<span class="k">for</span> <span class="n">t</span> <span class="ow">in</span> <span class="p">(</span><span class="nb">list</span><span class="p">(</span><span class="n">converted_positional_hints</span><span class="p">)</span> <span class="o">+</span>
<span class="nb">list</span><span class="p">(</span><span class="n">converted_keyword_hints</span><span class="o">.</span><span class="n">values</span><span class="p">())):</span>
<span class="n">validate_composite_type_param</span><span class="p">(</span>
<span class="n">t</span><span class="p">,</span> <span class="n">error_msg_prefix</span><span class="o">=</span><span class="s1">&#39;All type hint arguments&#39;</span><span class="p">)</span>
<span class="n">th</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="s1">&#39;_type_hints&#39;</span><span class="p">,</span> <span class="n">IOTypeHints</span><span class="o">.</span><span class="n">empty</span><span class="p">())</span><span class="o">.</span><span class="n">with_input_types</span><span class="p">(</span>
<span class="o">*</span><span class="n">converted_positional_hints</span><span class="p">,</span> <span class="o">**</span><span class="n">converted_keyword_hints</span><span class="p">)</span>
<span class="n">f</span><span class="o">.</span><span class="n">_type_hints</span> <span class="o">=</span> <span class="n">th</span> <span class="c1"># pylint: disable=protected-access</span>
<span class="k">return</span> <span class="n">f</span>
<span class="k">return</span> <span class="n">annotate_input_types</span></div>
<div class="viewcode-block" id="with_output_types"><a class="viewcode-back" href="../../../apache_beam.typehints.decorators.html#apache_beam.typehints.decorators.with_output_types">[docs]</a><span class="k">def</span> <span class="nf">with_output_types</span><span class="p">(</span><span class="o">*</span><span class="n">return_type_hint</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
<span class="c1"># type: (*Any, **Any) -&gt; Callable[[T], T]</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;A decorator that type-checks defined type-hints for return values(s).</span>
<span class="sd"> This decorator will type-check the return value(s) of the decorated function.</span>
<span class="sd"> Only a single type-hint is accepted to specify the return type of the return</span>
<span class="sd"> value. If the function to be decorated has multiple return values, then one</span>
<span class="sd"> should use: ``Tuple[type_1, type_2]`` to annotate the types of the return</span>
<span class="sd"> values.</span>
<span class="sd"> If the ultimate return value for the function violates the specified type-hint</span>
<span class="sd"> a :class:`TypeCheckError` will be raised detailing the type-constraint</span>
<span class="sd"> violation.</span>
<span class="sd"> This decorator is intended to be used like:</span>
<span class="sd"> .. testcode::</span>
<span class="sd"> from apache_beam.typehints import with_output_types</span>
<span class="sd"> from apache_beam.typehints import Set</span>
<span class="sd"> class Coordinate(object):</span>
<span class="sd"> def __init__(self, x, y):</span>
<span class="sd"> self.x = x</span>
<span class="sd"> self.y = y</span>
<span class="sd"> @with_output_types(Set[Coordinate])</span>
<span class="sd"> def parse_ints(ints):</span>
<span class="sd"> return {Coordinate(i, i) for i in ints}</span>
<span class="sd"> Or with a simple type-hint:</span>
<span class="sd"> .. testcode::</span>
<span class="sd"> from apache_beam.typehints import with_output_types</span>
<span class="sd"> @with_output_types(bool)</span>
<span class="sd"> def negate(p):</span>
<span class="sd"> return not p if p else p</span>
<span class="sd"> Args:</span>
<span class="sd"> *return_type_hint: A type-hint specifying the proper return type of the</span>
<span class="sd"> function. This argument should either be a built-in Python type or an</span>
<span class="sd"> instance of a :class:`~apache_beam.typehints.typehints.TypeConstraint`</span>
<span class="sd"> created by &#39;indexing&#39; a</span>
<span class="sd"> :class:`~apache_beam.typehints.typehints.CompositeTypeHint`.</span>
<span class="sd"> **kwargs: Not used.</span>
<span class="sd"> Raises:</span>
<span class="sd"> :class:`ValueError`: If any kwarg parameters are passed in,</span>
<span class="sd"> or the length of **return_type_hint** is greater than ``1``. Or if the</span>
<span class="sd"> inner wrapper function isn&#39;t passed a function object.</span>
<span class="sd"> :class:`TypeCheckError`: If the **return_type_hint** object is</span>
<span class="sd"> in invalid type-hint.</span>
<span class="sd"> Returns:</span>
<span class="sd"> The original function decorated such that it enforces type-hint constraints</span>
<span class="sd"> for all return values.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">if</span> <span class="n">kwargs</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
<span class="s2">&quot;All arguments for the &#39;returns&#39; decorator must be &quot;</span>
<span class="s2">&quot;positional arguments.&quot;</span><span class="p">)</span>
<span class="k">if</span> <span class="nb">len</span><span class="p">(</span><span class="n">return_type_hint</span><span class="p">)</span> <span class="o">!=</span> <span class="mi">1</span><span class="p">:</span>
<span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span>
<span class="s2">&quot;&#39;returns&#39; accepts only a single positional argument. In &quot;</span>
<span class="s2">&quot;order to specify multiple return types, use the &#39;Tuple&#39; &quot;</span>
<span class="s2">&quot;type-hint.&quot;</span><span class="p">)</span>
<span class="n">return_type_hint</span> <span class="o">=</span> <span class="n">native_type_compatibility</span><span class="o">.</span><span class="n">convert_to_beam_type</span><span class="p">(</span>
<span class="n">return_type_hint</span><span class="p">[</span><span class="mi">0</span><span class="p">])</span>
<span class="n">validate_composite_type_param</span><span class="p">(</span>
<span class="n">return_type_hint</span><span class="p">,</span> <span class="n">error_msg_prefix</span><span class="o">=</span><span class="s1">&#39;All type hint arguments&#39;</span><span class="p">)</span>
<span class="k">def</span> <span class="nf">annotate_output_types</span><span class="p">(</span><span class="n">f</span><span class="p">):</span>
<span class="n">th</span> <span class="o">=</span> <span class="nb">getattr</span><span class="p">(</span><span class="n">f</span><span class="p">,</span> <span class="s1">&#39;_type_hints&#39;</span><span class="p">,</span> <span class="n">IOTypeHints</span><span class="o">.</span><span class="n">empty</span><span class="p">())</span>
<span class="n">f</span><span class="o">.</span><span class="n">_type_hints</span> <span class="o">=</span> <span class="n">th</span><span class="o">.</span><span class="n">with_output_types</span><span class="p">(</span><span class="n">return_type_hint</span><span class="p">)</span> <span class="c1"># pylint: disable=protected-access</span>
<span class="k">return</span> <span class="n">f</span>
<span class="k">return</span> <span class="n">annotate_output_types</span></div>
<span class="k">def</span> <span class="nf">_check_instance_type</span><span class="p">(</span>
<span class="n">type_constraint</span><span class="p">,</span> <span class="n">instance</span><span class="p">,</span> <span class="n">var_name</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span> <span class="n">verbose</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;A helper function to report type-hint constraint violations.</span>
<span class="sd"> Args:</span>
<span class="sd"> type_constraint: An instance of a &#39;TypeConstraint&#39; or a built-in Python</span>
<span class="sd"> type.</span>
<span class="sd"> instance: The candidate object which will be checked by to satisfy</span>
<span class="sd"> &#39;type_constraint&#39;.</span>
<span class="sd"> var_name: If &#39;instance&#39; is an argument, then the actual name for the</span>
<span class="sd"> parameter in the original function definition.</span>
<span class="sd"> Raises:</span>
<span class="sd"> TypeCheckError: If &#39;instance&#39; fails to meet the type-constraint of</span>
<span class="sd"> &#39;type_constraint&#39;.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="n">hint_type</span> <span class="o">=</span> <span class="p">(</span>
<span class="s2">&quot;argument: &#39;</span><span class="si">%s</span><span class="s2">&#39;&quot;</span> <span class="o">%</span> <span class="n">var_name</span> <span class="k">if</span> <span class="n">var_name</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span> <span class="k">else</span> <span class="s1">&#39;return type&#39;</span><span class="p">)</span>
<span class="k">try</span><span class="p">:</span>
<span class="n">check_constraint</span><span class="p">(</span><span class="n">type_constraint</span><span class="p">,</span> <span class="n">instance</span><span class="p">)</span>
<span class="k">except</span> <span class="n">SimpleTypeHintError</span><span class="p">:</span>
<span class="k">if</span> <span class="n">verbose</span><span class="p">:</span>
<span class="n">verbose_instance</span> <span class="o">=</span> <span class="s1">&#39;</span><span class="si">%s</span><span class="s1">, &#39;</span> <span class="o">%</span> <span class="n">instance</span>
<span class="k">else</span><span class="p">:</span>
<span class="n">verbose_instance</span> <span class="o">=</span> <span class="s1">&#39;&#39;</span>
<span class="k">raise</span> <span class="n">TypeCheckError</span><span class="p">(</span>
<span class="s1">&#39;Type-hint for </span><span class="si">%s</span><span class="s1"> violated. Expected an &#39;</span>
<span class="s1">&#39;instance of </span><span class="si">%s</span><span class="s1">, instead found </span><span class="si">%s</span><span class="s1">an instance of </span><span class="si">%s</span><span class="s1">.&#39;</span> <span class="o">%</span>
<span class="p">(</span><span class="n">hint_type</span><span class="p">,</span> <span class="n">type_constraint</span><span class="p">,</span> <span class="n">verbose_instance</span><span class="p">,</span> <span class="nb">type</span><span class="p">(</span><span class="n">instance</span><span class="p">)))</span>
<span class="k">except</span> <span class="n">CompositeTypeHintError</span> <span class="k">as</span> <span class="n">e</span><span class="p">:</span>
<span class="k">raise</span> <span class="n">TypeCheckError</span><span class="p">(</span><span class="s1">&#39;Type-hint for </span><span class="si">%s</span><span class="s1"> violated: </span><span class="si">%s</span><span class="s1">&#39;</span> <span class="o">%</span> <span class="p">(</span><span class="n">hint_type</span><span class="p">,</span> <span class="n">e</span><span class="p">))</span>
<span class="k">def</span> <span class="nf">_interleave_type_check</span><span class="p">(</span><span class="n">type_constraint</span><span class="p">,</span> <span class="n">var_name</span><span class="o">=</span><span class="kc">None</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;Lazily type-check the type-hint for a lazily generated sequence type.</span>
<span class="sd"> This function can be applied as a decorator or called manually in a curried</span>
<span class="sd"> manner:</span>
<span class="sd"> * @_interleave_type_check(List[int])</span>
<span class="sd"> def gen():</span>
<span class="sd"> yield 5</span>
<span class="sd"> or</span>
<span class="sd"> * gen = _interleave_type_check(Tuple[int, int], &#39;coord_gen&#39;)(gen)</span>
<span class="sd"> As a result, all type-checking for the passed generator will occur at &#39;yield&#39;</span>
<span class="sd"> time. This way, we avoid having to depleat the generator in order to</span>
<span class="sd"> type-check it.</span>
<span class="sd"> Args:</span>
<span class="sd"> type_constraint: An instance of a TypeConstraint. The output yielded of</span>
<span class="sd"> &#39;gen&#39; will be type-checked according to this type constraint.</span>
<span class="sd"> var_name: The variable name binded to &#39;gen&#39; if type-checking a function</span>
<span class="sd"> argument. Used solely for templating in error message generation.</span>
<span class="sd"> Returns:</span>
<span class="sd"> A function which takes a generator as an argument and returns a wrapped</span>
<span class="sd"> version of the generator that interleaves type-checking at &#39;yield&#39;</span>
<span class="sd"> iteration. If the generator received is already wrapped, then it is simply</span>
<span class="sd"> returned to avoid nested wrapping.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">def</span> <span class="nf">wrapper</span><span class="p">(</span><span class="n">gen</span><span class="p">):</span>
<span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">gen</span><span class="p">,</span> <span class="n">GeneratorWrapper</span><span class="p">):</span>
<span class="k">return</span> <span class="n">gen</span>
<span class="k">return</span> <span class="n">GeneratorWrapper</span><span class="p">(</span>
<span class="n">gen</span><span class="p">,</span> <span class="k">lambda</span> <span class="n">x</span><span class="p">:</span> <span class="n">_check_instance_type</span><span class="p">(</span><span class="n">type_constraint</span><span class="p">,</span> <span class="n">x</span><span class="p">,</span> <span class="n">var_name</span><span class="p">))</span>
<span class="k">return</span> <span class="n">wrapper</span>
<span class="k">class</span> <span class="nc">GeneratorWrapper</span><span class="p">(</span><span class="nb">object</span><span class="p">):</span>
<span class="w"> </span><span class="sd">&quot;&quot;&quot;A wrapper around a generator, allows execution of a callback per yield.</span>
<span class="sd"> Additionally, wrapping a generator with this class allows one to assign</span>
<span class="sd"> arbitary attributes to a generator object just as with a function object.</span>
<span class="sd"> Attributes:</span>
<span class="sd"> internal_gen: A instance of a generator object. As part of &#39;step&#39; of the</span>
<span class="sd"> generator, the yielded object will be passed to &#39;interleave_func&#39;.</span>
<span class="sd"> interleave_func: A callback accepting a single argument. This function will</span>
<span class="sd"> be called with the result of each yielded &#39;step&#39; in the internal</span>
<span class="sd"> generator.</span>
<span class="sd"> &quot;&quot;&quot;</span>
<span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">gen</span><span class="p">,</span> <span class="n">interleave_func</span><span class="p">):</span>
<span class="bp">self</span><span class="o">.</span><span class="n">internal_gen</span> <span class="o">=</span> <span class="n">gen</span>
<span class="bp">self</span><span class="o">.</span><span class="n">interleave_func</span> <span class="o">=</span> <span class="n">interleave_func</span>
<span class="k">def</span> <span class="fm">__getattr__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">attr</span><span class="p">):</span>
<span class="c1"># TODO(laolu): May also want to intercept &#39;send&#39; in the future if we move to</span>
<span class="c1"># a GeneratorHint with 3 type-params:</span>
<span class="c1"># * Generator[send_type, return_type, yield_type]</span>
<span class="k">if</span> <span class="n">attr</span> <span class="o">==</span> <span class="s1">&#39;__next__&#39;</span><span class="p">:</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="fm">__next__</span><span class="p">()</span>
<span class="k">elif</span> <span class="n">attr</span> <span class="o">==</span> <span class="s1">&#39;__iter__&#39;</span><span class="p">:</span>
<span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="fm">__iter__</span><span class="p">()</span>
<span class="k">return</span> <span class="nb">getattr</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">internal_gen</span><span class="p">,</span> <span class="n">attr</span><span class="p">)</span>
<span class="k">def</span> <span class="fm">__next__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="n">next_val</span> <span class="o">=</span> <span class="nb">next</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">internal_gen</span><span class="p">)</span>
<span class="bp">self</span><span class="o">.</span><span class="n">interleave_func</span><span class="p">(</span><span class="n">next_val</span><span class="p">)</span>
<span class="k">return</span> <span class="n">next_val</span>
<span class="nb">next</span> <span class="o">=</span> <span class="fm">__next__</span>
<span class="k">def</span> <span class="fm">__iter__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
<span class="k">for</span> <span class="n">x</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">internal_gen</span><span class="p">:</span>
<span class="bp">self</span><span class="o">.</span><span class="n">interleave_func</span><span class="p">(</span><span class="n">x</span><span class="p">)</span>
<span class="k">yield</span> <span class="n">x</span>
</pre></div>
</div>
</div>
<footer>
<hr/>
<div role="contentinfo">
<p>
&copy; Copyright
</p>
</div>
Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a> provided by <a href="https://readthedocs.org">Read the Docs</a>.
</footer>
</div>
</div>
</section>
</div>
<script type="text/javascript">
jQuery(function () {
SphinxRtdTheme.Navigation.enable(true);
});
</script>
</body>
</html>