SourceKitten Swift items extractor

SwiftItemsExtractor.py

import subprocess
import json

class SwiftItemsExtractor:

    SOURCEKIT_SUBSCTRUCTURE_KEY = "key.substructure"
    SOURCEKIT_OFFSET_KEY = "key.offset"
    SOURCEKIT_LENGTH_KEY = "key.length"
    SOURCEKIT_COMMENT_KEY = "key.doc.comment"
    SOURCEKIT_NAME_KEY = "key.name"
    SOURCEKIT_KIND_KEY = "key.kind"
    SOURCEKIT_CLASS_TYPE_VALUE = "source.lang.swift.decl.class"
    SOURCEKIT_METHOD_TYPE_VALUE = "source.lang.swift.decl.function.method.instance"

    def swift_struct(self, file_path):
        sourcekitten_json = subprocess.Popen("sourcekitten structure --file %s" % file_path, shell=True, stdout=subprocess.PIPE).stdout.read()
        try:
          return json.loads(sourcekitten_json)
        except:
          return {}
    
    def swift_code_struct(self, code_block):
        sourcekitten_json = subprocess.Popen("sourcekitten structure --text \"%s\"" % code_block, shell=True, stdout=subprocess.PIPE).stdout.read()
        try:
          return json.loads(sourcekitten_json)
        except Exception as e:
            print(e)
            return {}


    def get_all_methods_from_sourcekitten(self, file_structure, kind_of_items):
       result_items = []

       result_type = ""
       if kind_of_items == 'class':
          result_type = SwiftMethodsExtractor.SOURCEKIT_CLASS_TYPE_VALUE
       elif kind_of_items == 'method':
          result_type = SwiftMethodsExtractor.SOURCEKIT_METHOD_TYPE_VALUE

       result_items.extend(self.get_methods_from_sourcekitten(file_structure, result_type))

       if SwiftMethodsExtractor.SOURCEKIT_SUBSCTRUCTURE_KEY not in file_structure:
          return result_items

       for item in file_structure[SwiftMethodsExtractor.SOURCEKIT_SUBSCTRUCTURE_KEY]:
          result_items.extend(self.get_methods_from_sourcekitten(item, result_type))
          result_items.extend(self.get_all_methods_from_sourcekitten(item, kind_of_items))
       return result_items


    def get_methods_from_sourcekitten(self, file_structure, kind):
        if SwiftMethodsExtractor.SOURCEKIT_SUBSCTRUCTURE_KEY not in file_structure:
          return []

        return [
          { 
            "start":item[SwiftMethodsExtractor.SOURCEKIT_OFFSET_KEY],
            "length":item[SwiftMethodsExtractor.SOURCEKIT_LENGTH_KEY],
            "doc":item[SwiftMethodsExtractor.SOURCEKIT_COMMENT_KEY] if SwiftMethodsExtractor.SOURCEKIT_COMMENT_KEY in item else None,
            "name":item[SwiftMethodsExtractor.SOURCEKIT_NAME_KEY]
          } 
          for item 
          in file_structure[SwiftMethodsExtractor.SOURCEKIT_SUBSCTRUCTURE_KEY] 
          if item[SwiftMethodsExtractor.SOURCEKIT_KIND_KEY] == kind
        ]


    def process_code_block(self, file_body, meta):

      result_items = []

      if meta["doc"]:
        result_items.append(re.sub(r'^', r'//  ', meta["doc"], flags=re.M).encode('utf-8'))

      body = file_body[meta["start"]:meta["start"]+meta["length"]]
      
      last_line = body.split('\n')[-1]
      intent_regex = re.compile('^%s' % last_line.replace('}', ''), flags=re.M)
      result_items.append(re.sub(intent_regex, r'', body))

      return "\n".join(result_items)

    def get_methods_from_swift_file(self, file_path, kind_of_items):
        swift_structs = self.swift_struct(file_path).values()
        if not swift_structs:
          return []
        
        print(swift_structs)
        swift_struct = swift_structs[0]

        methods_meta = self.get_all_methods_from_sourcekitten(swift_struct, kind_of_items)
        file_data = open(file_path, "rb+").read()
        return [
          {
            "body": self.process_code_block(file_data, meta),
            "comments": meta["doc"],
            "file_path": file_path
          }
          for meta 
          in methods_meta
        ]