#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

import argparse

import glob

import io

import itertools

import os

import re

import docstring_parser

import yaml

from apache_beam.portability.api import schema_pb2

from apache_beam.typehints import schemas

from apache_beam.utils import subprocess_server

from apache_beam.utils.python_callable import PythonCallableWithSource

from apache_beam.version import __version__ as beam_version

from apache_beam.yaml import json_utils

from apache_beam.yaml import yaml_provider

from apache_beam.yaml import yaml_utils

from apache_beam.yaml.yaml_errors import ErrorHandlingConfig

def _singular(name):

# Simply removing an 's' (or 'es', or 'ies', ...) may result in surprising

# manglings. Better to play it safe and leave a correctly-spelled plural

# than a botched singular in our examples configs.

return {

'args': 'arg',

'attributes': 'attribute',

'elements': 'element',

'fields': 'field',

}.get(name, name)

def _fake_value(name, beam_type):

type_info = beam_type.WhichOneof("type_info")

if type_info == "atomic_type":

if beam_type.atomic_type == schema_pb2.STRING:

return f'"{name}"'

elif beam_type.atomic_type == schema_pb2.BOOLEAN:

return "true|false"

else:

return name

elif type_info == "array_type":

return [

_fake_value(_singular(name), beam_type.array_type.element_type),

_fake_value(_singular(name), beam_type.array_type.element_type),

'...'

]

elif type_info == "iterable_type":

return [

_fake_value(_singular(name), beam_type.iterable_type.element_type),

_fake_value(_singular(name), beam_type.iterable_type.element_type),

'...'

]

elif type_info == "map_type":

if beam_type.map_type.key_type.atomic_type == schema_pb2.STRING:

return {

'a': _fake_value(name + '_value_a', beam_type.map_type.value_type),

'b': _fake_value(name + '_value_b', beam_type.map_type.value_type),

'c': '...',

}

else:

return {

_fake_value(name + '_key', beam_type.map_type.key_type): _fake_value(

name + '_value', beam_type.map_type.value_type)

}

elif type_info == "row_type":

return _fake_row(beam_type.row_type.schema)

elif type_info == "logical_type":

return name

else:

raise ValueError(f"Unrecognized type_info: {type_info!r}")

def _fake_row(schema):

if schema is None:

return '...'

return {f.name: _fake_value(f.name, f.type) for f in schema.fields}

def pretty_example(provider, t, base_t=None):

spec = {'type': base_t or t}

try:

requires_inputs = provider.requires_inputs(t, {})

except Exception:

requires_inputs = False

if requires_inputs:

spec['input'] = '...'

config_schema = provider.config_schema(t)

if config_schema is None or config_schema.fields:

spec['config'] = _fake_row(config_schema)

s = yaml.dump(spec, sort_keys=False)

return s.replace("'", "")

def config_docs(schema):

if schema is None:

return ''

elif not schema.fields:

return 'No configuration parameters.'

def pretty_type(beam_type):

type_info = beam_type.WhichOneof("type_info")

if type_info == "atomic_type":

return schema_pb2.AtomicType.Name(beam_type.atomic_type).lower()

elif type_info == "array_type":

return f'Array[{pretty_type(beam_type.array_type.element_type)}]'

elif type_info == "iterable_type":

return f'Iterable[{pretty_type(beam_type.iterable_type.element_type)}]'

elif type_info == "map_type":

return (

f'Map[{pretty_type(beam_type.map_type.key_type)}, '

f'{pretty_type(beam_type.map_type.value_type)}]')

elif type_info == "row_type":

return 'Row'

else:

return '?'

def maybe_row_parameters(t):

if t.WhichOneof("type_info") == "row_type":

return indent('\n\nRow fields:\n\n' + config_docs(t.row_type.schema), 4)

else:

return ''

def maybe_optional(t):

return " (Optional)" if t.nullable else ""

def normalize_error_handling(f):

doc = docstring_parser.parse(

ErrorHandlingConfig.__doc__, docstring_parser.DocstringStyle.GOOGLE)

if f.name == "error_handling":

f = schema_pb2.Field(

name="error_handling",

type=schema_pb2.FieldType(

row_type=schema_pb2.RowType(

schema=schema_pb2.Schema(

fields=[

schemas.schema_field(

param.arg_name,

PythonCallableWithSource.load_from_expression(

param.type_name),

param.description) for param in doc.params

])),

nullable=True),

description=f.description or doc.short_description)

return f

def lines():

for f in schema.fields:

f = normalize_error_handling(f)

yield ''.join([

f'**{f.name}** `{pretty_type(f.type)}`',

maybe_optional(f.type),

indent(': ' + f.description if f.description else '', 2),

maybe_row_parameters(f.type),

])

return '\n\n'.join('*' + indent(line, 2) for line in lines()).strip()

def indent(lines, size):

return '\n'.join(' ' * size + line for line in lines.split('\n'))

def longest(func, xs):

return max([func(x) or '' for x in xs], key=len)

def io_grouping_key(transform_name):

"""Place reads and writes next to each other, after all other transforms."""

if transform_name.startswith('ReadFrom'):

return 1, transform_name[8:], 0

elif transform_name.startswith('WriteTo'):

return 1, transform_name[7:], 1

else:

return 0, transform_name

# Exclude providers

SKIP = {}

def add_transform_links(transform, description, provider_list):

"""

Convert references of Providers to urls that link to their respective pages.

For example,

"Some description talking about MyTransform."

would be converted to

"Some description talking about <a href="#mytransform">MyTransform</a>"

meanwhile::

type: MyTransform

config:

...

Would remain unchanged.

Avoid self-linking within a Transform page.

"""

for p in provider_list:

# Match all instances of built-in transforms within the description

# excluding the transform whose description is currently being evaluated.

# Match the entire word boundary so that partial matches do not count.

# (i.e. OtherTransform should not match Transform)

description = re.sub(

rf"(?<!type: )\b(?!{transform}\b)\b{p}\b",

f'<a href="#{p.lower()}">{p}</a>',

description or '')

return description

def transform_docs(transform_base, transforms, providers, extra_docs=''):

return '\n'.join([

f'## {transform_base}',

'',

longest(

lambda t: longest(

lambda p: add_transform_links(

t, p.description(t), providers.keys()), providers[t]),

transforms).replace('::\n', '\n\n :::yaml\n'),

'',

extra_docs,

'',

'### Configuration',

'',

longest(

lambda t: longest(

lambda p: config_docs(p.config_schema(t)), providers[t]),

transforms),

'',

'### Usage',

'',

' :::yaml',

'',

indent(

longest(

lambda t: longest(

lambda p: pretty_example(p, t, transform_base), providers[t]),

transforms),

4),

])

def create_index(include, exclude, options):

with subprocess_server.SubprocessServer.cache_subprocesses():

json_config_schemas = []

markdown_out = io.StringIO()

if options.provider_source_file:

providers = yaml_provider.merge_providers(

yaml_provider.load_providers(

yaml_utils.locate_data_file(options.provider_source_file)))

else:

providers = yaml_provider.standard_providers()

for transform_base, transforms in itertools.groupby(

sorted(providers.keys(), key=io_grouping_key),

key=lambda s: s.split('-')[0]):

transforms = list(transforms)

if include(transform_base) and not exclude(transform_base):

print(transform_base)

if options.markdown_file or options.html_file:

if '-' in transforms[0]:

extra_docs = 'Supported languages: ' + ', '.join(

t.split('-')[-1] for t in sorted(transforms)) + '.'

else:

extra_docs = ''

markdown_out.write(

transform_docs(transform_base, transforms, providers, extra_docs))

markdown_out.write('\n\n')

if options.schema_file:

for transform in transforms:

schema = providers[transform][0].config_schema(transform)

if schema:

json_config_schemas.append({

'if': {

'properties': {

'type': {

'const': transform

}

}

},

'then': {

'properties': {

'config': {

'type': 'object',

'properties': {

'__line__': {

'type': 'integer'

},

'__uuid__': {},

**{

f.name: #

json_utils.beam_type_to_json_type(f.type)

for f in schema.fields

}

},

'additionalProperties': False,

}

}

}

})

return json_config_schemas, markdown_out.getvalue()

def create_examples_markdown():

markdown_out = io.StringIO()

base = os.path.join(os.path.dirname(__file__), 'examples')

section = last_section = ''

for path in sorted(glob.glob(os.path.join(base, '**', '*.yaml'),

recursive=True),

key=lambda path: (path.count(os.sep), path)):

short_path = path[len(base):].replace('transforms', '').strip(os.sep)

def to_title(path):

base, _ = os.path.splitext(path)

nice = base.replace('_', ' ').replace(os.sep, ' ').title()

# These acronyms should be upper, not title.

nice = re.sub(r'\bMl\b', 'ML', nice)

nice = re.sub(r'\bIo\b', 'IO', nice)

return nice

def clean_yaml(content):

content = re.sub(

'# Licensed to the Apache Software Foundation.*'

'# limitations under the License.',

'',

content,

flags=re.MULTILINE | re.DOTALL)

content = re.sub('# coding=.*', '', content)

return content

def split_header(yaml):

lines = yaml.split('\n')

ix = 0 # make lint happy

for ix, line in enumerate(lines):

if not line.strip():

continue

if not line.startswith('#'):

break

return (

'\n'.join([line[1:].strip() for line in lines[:ix]]),

'\n'.join(lines[ix:]))

if os.sep in short_path:

section = to_title(short_path.split(os.sep)[0])

if section != last_section:

markdown_out.write(f'# {section}\n\n')

last_section = section

title = to_title(short_path)[len(section):]

markdown_out.write(f'## {title}\n\n')

with open(path) as fin:

content = fin.read()

header, body = split_header(clean_yaml(content))

markdown_out.write(header)

markdown_out.write('\n\n :::yaml\n\n')

markdown_out.write(' ' + body.replace('\n', '\n '))

markdown_out.write('\n')

return markdown_out.getvalue()

def markdown_to_html(title, markdown_content, header=''):

import markdown

import markdown.extensions.toc

import pygments.formatters

md = markdown.Markdown(

extensions=[

markdown.extensions.toc.TocExtension(toc_depth=2),

'codehilite',

])

pygments_style = pygments.formatters.HtmlFormatter().get_style_defs(

'.codehilite')

extra_style = '''

* {

box-sizing: border-box;

}

body {

font-family: 'Roboto', sans-serif;

font-weight: normal;

color: #404040;

background: #edf0f2;

}

.body-for-nav {

background: #fcfcfc;

}

.grid-for-nav {

width: 100%;

}

.nav-side {

position: fixed;

top: 0;

left: 0;

width: 300px;

height: 100%;

padding-bottom: 2em;

color: #9b9b9b;

background: #343131;

}

.nav-header {

display: block;

width: 300px;

padding: 1em;

background-color: #2980B9;

text-align: center;

color: #fcfcfc;

}

.nav-header a {

color: #fcfcfc;

font-weight: bold;

display: inline-block;

padding: 4px 6px;

margin-bottom: 1em;

text-decoration:none;

}

.nav-header>div.version {

margin-top: -.5em;

margin-bottom: 1em;

font-weight: normal;

color: rgba(255, 255, 255, 0.3);

}

.toc {

width: 300px;

text-align: left;

overflow-y: auto;

max-height: calc(100% - 4.3em);

scrollbar-width: thin;

scrollbar-color: #9b9b9b #343131;

}

.toc ul {

margin: 0;

padding: 0;

list-style: none;

}

.toc li {

border-bottom: 1px solid #4e4a4a;

margin-left: 1em;

}

.toc a {

display: block;

line-height: 36px;

font-size: 90%;

color: #d9d9d9;

padding: .1em 0.6em;

text-decoration: none;

transition: background-color 0.3s ease, color 0.3s ease;

}

.toc a:hover {

background-color: #4e4a4a;

color: #ffffff;

}

.transform-content-wrap {

margin-left: 300px;

background: #fcfcfc;

}

.transform-content {

padding: 1.5em 3em;

margin: 20px;

padding-bottom: 2em;

}

.transform-content li::marker {

display: inline-block;

width: 0.5em;

}

.transform-content h1 {

font-size: 40px;

}

.transform-content ul {

margin-left: 0.75em;

text-align: left;

list-style-type: disc;

}

hr {

color: gray;

display: block;

height: 1px;

border: 0;

border-top: 1px solid #e1e4e5;

margin-bottom: 3em;

margin-top: 3em;

padding: 0;

}

.codehilite {

background: #f5f5f5;

border: 1px solid #ccc;

border-radius: 4px;

padding: 0.2em 1em;

overflow: auto;

font-family: monospace;

font-size: 14px;

line-height: 1.5;

}

p code, li code {

white-space: nowrap;

max-width: 100%;

background: #fff;

border: solid 1px #e1e4e5;

padding: 0 5px;

font-family: monospace;

color: #404040;

font-weight: bold;

padding: 2px 5px;

}

'''

html = md.convert(markdown_content)

return f'''

<html>

<head>

<meta charset="UTF-8">

<title>{title}</title>

<style>

{pygments_style}

{extra_style}

</style>

</head>

<body class="body-for-nav">

<div class="grid-for-nav">

<nav class="nav-side">

<div class="nav-header">

<a href=#>{title}</a>

<div class="version">

{beam_version}

</div>

</div>

{getattr(md, 'toc')}

</nav>

<section class="transform-content-wrap">

<div class="transform-content">

<h1>{title}</h1>

{header}

{html.replace('<h2', '<hr><h2')}

</div>

</section>

</div>

</body>

</html>

'''

def main():

parser = argparse.ArgumentParser()

parser.add_argument('--examples_file')

parser.add_argument('--markdown_file')

parser.add_argument('--html_file')

parser.add_argument('--schema_file')

parser.add_argument('--include', default='.*')

parser.add_argument('--exclude', default='')

parser.add_argument(

"--provider_source_file",

help="Path to a YAML file containing custom providers. "

"If not provided, uses standard Beam providers.",

)

options = parser.parse_args()

include = re.compile(options.include).match

exclude = (

re.compile(options.exclude).match

if options.exclude else lambda x: x in SKIP)

json_config_schemas, markdown_content = create_index(

include, exclude, options)

if options.schema_file:

with open(options.schema_file, 'w') as fout:

yaml.dump(json_config_schemas, fout, sort_keys=False)

if options.markdown_file:

with open(options.markdown_file, 'w') as fout:

fout.write(markdown_content)

if options.html_file:

with open(options.html_file, 'w') as html_out:

html_out.write(

markdown_to_html('Beam YAML Transform Index', markdown_content))

if options.examples_file:

with open(options.examples_file, 'w') as html_out:

html_out.write(

markdown_to_html(

'Beam YAML Examples',

create_examples_markdown(),

header='''

<p>Example pipelines using the

<a href="https://beam.apache.org/documentation/sdks/yaml/">

Beam YAML API</a>.

These examples can also be found on

<a href="https://github.com/apache/beam/tree/master/sdks/'''

'''python/apache_beam/yaml/examples">github</a>.

</p>

<p>Note: These examples below are automatically tested for

correctness and may be used as a starting point for your own

pipelines.</p>

'''))

if __name__ == '__main__':

main()