walk.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import os
import re
import sys


def needsTag(idx, lines):
  idx = idx - 1  # Skip "public void foo()" line

  while idx >= 0:
    if '@PublicApi' in lines[idx]:
      return False
    if '@VisibleForTesting' in lines[idx]:
      return False
    if '@RestrictTo' in lines[idx]:
      return False
    if '@Override' in lines[idx]:
      return False
    if '@' not in lines[idx]:
      return True

    idx = idx - 1

  return False


def walk_dir(root_dir):
  root_dir = os.path.abspath(root_dir)

  for item in os.listdir(root_dir):
    full_path = os.path.join(root_dir, item)
    if os.path.isdir(full_path):
      walk_dir(full_path)
    elif full_path.endswith('.java'):
      contents = []
      modified = False

      with open(full_path, 'r') as f:
        print 'Processing ' + full_path
        contents = f.readlines()
        skipUntil = None  # Used to skip elements that are indented further

        for (idx, line) in enumerate(contents):
          if skipUntil is not None:
            if line == skipUntil:
              skipUntil = None
            else:
              continue

          match = re.match('^( *)public (.*)$', line)
          if match and needsTag(idx, contents):
            contents[idx] = \
              '@RestrictTo(RestrictTo.Scope.LIBRARY_GROUP) ' \
              + line
            skipUntil = match.groups(1)[0] + '}\n'
            modified = True
          else:
            # Skip non-public classes
            match = re.match('^( *)class (.*)$', line)
            if match:
              skipUntil = match.groups(1)[0] + '}\n'

      if modified:
        with open(full_path, 'w') as f:
          contents = contents[0:16] \
            + ['import android.support.annotation.RestrictTo;\n'
               ] + contents[16:]
          f.writelines(contents)


def main():
  walk_dir(str(sys.argv[1]))


if __name__ == '__main__':
  main()