yaml2asyncapi

#!/usr/bin/env python3

import sys
import json
from typing import Literal
import yaml
try:
  from yaml import CLoader as Loader, CDumper as Dumper
except ImportError:
  from yaml import Loader, Dumper

from yaml.representer import SafeRepresenter

# for dumping literal strings as scalar literal blocks


class LiteralString(str):
  pass


def literal_presenter(dumper, data):
  if type(data) is tuple:
    data = ''.join(data)
  return dumper.represent_scalar('tag:yaml.org,2002:str', data, style='|')


yaml.add_representer(LiteralString, literal_presenter)


def type_to_schema(t):
  schema = {}
  description = t.get('description', '')

  if description != '':
    schema['description'] = LiteralString(description)

  # type
  typetype = t['type']

  # if we are given a list of types, handle first
  if type(typetype) is list:
    oneof = []
    for typeitem in typetype:
      if typeitem in refs:
        oneof.append({'$ref': refs[typeitem]})
      else:
        oneof.append({'type': typeitem})
    schema['oneOf'] = oneof
  elif typetype == 'int':
    schema['type'] = 'integer'
  elif typetype == 'bool':
    schema['type'] = 'boolean'
  elif typetype == 'object':
    schema['type'] = 'object'
  elif typetype == 'list':
    schema['type'] = 'array'
  elif 'list[' in typetype:
    schema['type'] = 'array'
    subtype = typetype.replace('list[', '')[:-1]
    if subtype in refs:
      schema['items'] = [{'$ref': refs[subtype]}]
    else:
      schema['items'] = [{'type': subtype}]
  elif typetype in refs:
    schema['$ref'] = refs[typetype]
  else:
    schema['type'] = typetype

  # parameters if there are some
  params = t.get('parameters', None)
  if params is not None:
    schema['properties'] = {}
    schema['required'] = []
    optional = False
    for paramname in params:
      p = params[paramname]
      if type(p) is dict:
        if 'optional' in p and p['optional'] == True:
          optional = True
        pschema = type_to_schema(p)
      elif type(p) is list:
        oneof = []
        for typeitem in p:
          if typeitem in refs:
            oneof.append({'$ref': refs[typeitem]})
          else:
            oneof.append({'type': typeitem})
        pschema = {'oneOf': oneof}

      else:
        pschema = {'const': p}

      if optional == False:
        schema['required'].append(paramname)

      schema['properties'][paramname] = pschema
      schema['additionalProperties'] = False

  # check for 'options'
  options = t.get('options', None)
  if options is not None:
    schema['enum'] = options

  # example
  example = t.get('example', '')
  if example != '':
    exampleObject = example
    if type(example) is str:
      try:
        exampleObject = json.loads(example.replace('`', ''))
      except:
        exampleObject = LiteralString(example)
    schema['examples'] = [exampleObject]

  return schema


# =============================
# BEGIN MAIN
# =============================
if len(sys.argv) == 1:
  print(f'''
USAGE:

{sys.argv[0]} yaml_file output_name

Will parse my simple API yaml file to AsyncAPI.
  ''')
  exit()

source = sys.argv[1]
target = source + '-asyncapi.yaml'
target = target.replace('.yaml-asyncapi.yaml', '-asyncapi.yaml')
target = target.replace('.yml-asyncapi.yaml', '-asyncapi.yaml')

if len(sys.argv) == 3:
  target = sys.argv[2]
  if target == source:
    print('ERROR: source file and target file have the same name')
    exit(1)

# read input yaml file
print(f'READING: {source}')
with open(source, 'r') as f:
  t = f.read()
  data = yaml.load(t, Loader=Loader)

print('PROCESSING...')

# include root level asyncapi data
output = data['asyncapi']
output.update({'components': {'schemas': {}, 'messages': {}}, 'channels': {}})

refs = {}

# create schemas from the types in two passes.
# pass 1 creates the type refs
for typename in data['types']:
  if typename not in refs:
    refs[typename] = f'#/components/schemas/{typename}'

# pass 2 actually creates the types
for typename in data['types']:
  t = data['types'][typename]
  # print(t)
  schema = type_to_schema(t)
  output['components']['schemas'][typename] = schema

# create pub/sub messages from the message / response payloads
for channel in ['remote', 'stagedisplay']:
  channelPath = '/'+channel
  summary = data['channels'][channel]['summary']
  description = data['channels'][channel]['description']
  output['channels'][channelPath] = {
      'publish': {
          'summary': summary,
          'description': f'Send messages to the API\n\n{description}',
          'message': {
              'oneOf': []
          }
      },
      'subscribe': {
          'summary': summary,
          'description': f'Receive messages from the API\n\n{description}',
          'message': {
              'oneOf': []
          }
      },
  }

  for actionName in data['channels'][channel]['actions']:
    action = data['channels'][channel]['actions'][actionName]
    # print(action)
    pubName = f'{actionName}__pub'
    subName = f'{actionName}__sub'
    messageNames = {'publish': pubName, 'subscribe': subName}
    messageTitles = {
        'publish': f'{actionName} (publish)',
        'subscribe': f'{actionName} (subscribe)',
    }

    for pubsub in ['publish', 'subscribe']:
      messageName = messageNames[pubsub]
      message = {}
      message['name'] = messageName             # machine readable
      message['title'] = messageTitles[pubsub]  # human readable

      if 'summary' in action:
        message['summary'] = action['summary']

      # generate descriptive text for this message
      desc = action.get('description', '')
      if desc is None:
        desc = ''

      if pubsub == 'subscribe' and 'message' in action:
        foreign = messageNames['publish']
        linkref = f'#message-{foreign}'
        linktext = messageTitles["publish"]
        desc += f'\n\n**RELATED COMMAND**: [{linktext}]({linkref})\n'
      if pubsub == 'subscribe' and 'message' not in action:
        desc = desc.strip() + '\n\nTHIS IS A SUBSCRIBE-ONLY MESSAGE'

      if pubsub == 'publish' and 'response' in action:
        foreign = messageNames['subscribe']
        linkref = f'#message-{foreign}'
        linktext = messageTitles["subscribe"]
        desc += f'\n\n**RELATED RESPONSE**: [{linktext}]({linkref})\n'
        # message['x-response'] = {'$ref': f'#/components/messages/{foreign}'}

      if pubsub == 'publish' and 'response' not in action:
        desc = desc.strip() + '\n\nNO RESPONSE'
      message['description'] = LiteralString(desc)

      key = 'response'
      if pubsub == 'publish':
        key = 'message'
      if key in action and action[key] != 'none':
        refstring = f'#/components/messages/{messageName}'
        refs[messageName] = refstring

        # handle payload ... could be a 'type' or a 'payload'
        payload = action[key].get('payload', None)
        if payload is not None:
          forschema = {'type': 'object', 'parameters': payload}
          message['payload'] = type_to_schema(forschema)
        else:
          payloadtype = action[key].get('type', None)
          if payloadtype in refs:
            message['payload'] = {'$ref': refs[payloadtype]}
          else:
            print('ERROR: message payload was empty and no type was specified')
            print('here is the offending action')
            print(json.dumps(action, indent=2))
            message['payload'] = {}

        # handle example
        if 'example' in action[key]:
          example = action[key]['example']
          if type(example) is dict:
            exampleObject = example
          elif type(example) is str:
            jsonexample = action[key]['example'].replace('`', '')
            try:
              exampleObject = json.loads(jsonexample)
            except:
              exampleObject = LiteralString(jsonexample)
          message['examples'] = [{'payload': exampleObject}]

        # add to messages component
        output['components']['messages'][messageName] = message

        # add ref to the channel/pubsub/messages
        output['channels'][channelPath][pubsub]['message']['oneOf'].append(
            {'$ref': refstring, }
        )

        # print('== DEBUG ==')
        # print(json.dumps(action, indent=2))
        # print(yaml.dump(message, sort_keys=False, indent=2, width=60, default_flow_style=False))
        # input('Hit enter to continue ... ')

print(f'WRITING: {target}')
o = yaml.dump(output, sort_keys=False, indent=2, width=120, default_flow_style=False)
with open(target, 'w') as of:
  of.write(o)
print('DONE')