text: # Copyright (c) 2020 The Neuropod Authors
  #
  # Licensed 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 inspect
  import re
  
  class Function(object):
      """
      Represents a function
      """
  
      def __init__(self, name, shortdoc="", params=[]):
          super(Function, self).__init__()
          self.name = name
          self.shortdoc = shortdoc
          self.params = params
  
  
  class Parameter(object):
      """
      Represents a parameter to a function
      """
  
      def __init__(self, name):
          super(Parameter, self).__init__()
          self.name = name
          self.docs_arr = []
          self.has_default = False
          self.default = None
  
      def add_doc_line(self, line):
          """
          Add a line of documentation
          """
          self.docs_arr.append(line)
  
      def get_clean_doc(self):
          """
          Get a clean version of the documentation for this parameter
          """
          # Strips common leading whitespace
          return inspect.cleandoc("\n".join(self.docs_arr))
  
      def set_default(self, value):
          """
          Set the default value for this parameter
          """
          self.has_default = True
          self.default = value
  
  
  def parse_docstring(f):
      """
      Parse the docstring of `f` and return a list of parameters
      """
      doc = inspect.cleandoc(f.__doc__)
  
      lines = doc.splitlines()
      params = []
      currentparam = None
      shortdoc = ""
  
      def flush_param_if_necessary():
          if currentparam:
              params.append(currentparam)
  
      # Loop through all the lines of the docstring
      for i, line in enumerate(lines):
          match = re.match(r':param\s+(\w+):\s+(.*)', line)
          if match:
              # This is a new param
              flush_param_if_necessary()
  
              currentparam = Parameter(name=match.group(1))
              currentparam.add_doc_line(match.group(2))
          elif i == 0:
              # If this is the first line and it's not a new parameter,
              # this is the shortdoc
              shortdoc = line
          elif currentparam is not None:
              # This is part of the documentation of the current parameter
              currentparam.add_doc_line(line)
  
      # Make sure we get the final parameter
      flush_param_if_necessary()
  
      # Check for default values
      params_required = []
      params_defaults = []
      for param in params:
          if param.name in f.neuropod_default_args:
              param.set_default(f.neuropod_default_args[param.name])
              params_defaults.append(param)
          else:
              params_required.append(param)
  
      return Function(name=f.__name__, shortdoc=shortdoc, params=params_required + params_defaults)
  
  def write_doc(filename, f):
      """
      Write the documentation for a Function `f` to `filename`
      """
      with open(filename, "w") as docfile:
  
          # Write the title and shortdoc
          docfile.write("## " + f.name + "\n")
          docfile.write(f.shortdoc + "\n")
  
          # Write out the signature
          docfile.write("```\n")
          docfile.write(f.name + "(\n")
          for param in f.params:
              if param.has_default:
                  docfile.write("    " + param.name + " = " + str(param.default) + ",\n")
              else:
                  docfile.write("    " + param.name + ",\n")
          docfile.write(")\n")
          docfile.write("```\n")
  
          # Write out the list of params with their docs
          docfile.write("### Params:\n")
          for param in f.params:
              docfile.write("#### " + param.name + "\n")
              if param.has_default:
                  docfile.write("*default: `" + str(param.default) + "`*\n\n")
  
              docfile.write(param.get_clean_doc() + "\n\n")
  
  def write_doc_for_packager(packager, filename):
      """
      Writes the documentation for a packager to a file
      """
      f = parse_docstring(packager)
      write_doc(filename, f)
  
  if __name__ == '__main__':
      import argparse
      import os
  
      from neuropod.packagers import create_tensorflow_neuropod, \
                                      create_pytorch_neuropod, \
                                      create_keras_neuropod, \
                                      create_torchscript_neuropod
  
      parser = argparse.ArgumentParser(description='Generate markdown documentation for the Neuropod packagers')
      parser.add_argument('out_dir', help='The output directory to write the docs to')
      args = parser.parse_args()
  
      packager_mapping = {
          "tensorflow.md": create_tensorflow_neuropod,
          "pytorch.md": create_pytorch_neuropod,
          "keras.md": create_keras_neuropod,
          "torchscript.md": create_torchscript_neuropod,
      }
  
      for filename, packager in packager_mapping.items():
          write_doc_for_packager(packager, os.path.join(args.out_dir, filename))
  