| # | |
| # 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() |