annotate source code with a copyright header

annotate_source

#!/usr/bin/env python
"""
Annotate source code with a notice (see NOTICE_TEXT).

An existing notice will be replaced, and if there is no notice
encountered then one will be inserted. Detection of the notice
is exceedingly simple: if any line without a comment is encountered, from
the top of the file, before the standard "separator" of a long string
of comment characters, then the notice will be inserted. Likewise, the
"end" of the notice is either the same separator used for the beginning or
a line that is not commented.

For example, in the following he notice will be inserted between the
second and third lines::

    #!/usr/bin/env python
    # hello
    import sys

In this file he notice will be inserted before the first line::

    '''
    Top of the file comment
    '''
    import logging


Finally, if the notice is already there then the entire notice will be
replaced with the current text::

    ############################################################
    # Copyright (C) 2099 Nobody
    # You cannot have this code. Ever. It's too cool.
    ############################################################
    import asyncio

"""
import argparse
from collections import deque
from glob import glob, fnmatch
import logging
import os
import re
import shutil
import sys

_log = logging.getLogger('annotate_source')
_h = logging.StreamHandler()
_h.setFormatter(logging.Formatter(
    fmt='%(asctime)s [%(levelname)s] %(message)s'))
_log.addHandler(_h)

NOTICE_TEXT = '''My Project Name (MPN) Copyright (c) 2018, by the
software owners: The Regents of the University of Wonderland, through
Down the Rabbit Hole, et al. All rights reserved.

Please see the files COPYRIGHT.txt and LICENSE.txt for full copyright and
license information, respectively. Both files are also available online
at the URL "https://github.com/MPN/myproject".'''


def modify_files(finder, modifier, **flags):
    while True:
        try:
            f = finder.get()
        except IndexError:
            break
        modifier.modify(f, **flags)


def print_files(finder):
    while True:
        try:
            f = finder.get()
        except IndexError:
            break
        print(f)


class FileFinder(object):
    def __init__(self, root: str, glob_pat=None):
        if not os.path.isdir(root):
            raise FileNotFoundError('Root directory "{}"'.format(root))
        glob_pat = ['*.py'] if glob_pat is None else glob_pat
        neg_pat, pos_pat = [], []
        for p in glob_pat:
            if not p:
                pass
            if p[0] == '~':
                neg_pat.append(p[1:])
                _log.info('Negative pattern: {}'.format(p[1:]))
            else:
                pos_pat.append(p)
                _log.info('Positive pattern: {}'.format(p[1:]))
        self._root = root
        self._q = deque()
        for pat in pos_pat:
            self._find(pat, neg_pat)

    def __len__(self):
        return len(self._q)

    def _find(self, glob_pat, neg_pat):
        pat = os.path.join(self._root, '**', glob_pat)
        if neg_pat:
            # need to check each file, to eliminate bad ones
            for fpath in glob(pat, recursive=True):
                f, ok = os.path.basename(fpath), True
                # eliminate any that match a negative pattern
                for np in neg_pat:
                    _log.debug('Match file {} to pattern {}'.format(
                        f, np))
                    if fnmatch.fnmatchcase(f, np):
                        ok = False
                        break
                if ok:
                    self._q.append(fpath)
        else:
            # just grab all files
            self._q.extend(glob(pat, recursive=True))

    def get(self) -> str:
        item = self._q.pop()
        return item


class FileModifier(object):
    comment_pfx = '#'
    comment_sep = comment_pfx * 78
    comment_minsep = comment_pfx * 10

    def __init__(self, text: str):
        lines = [l.strip() for l in text.split('\n')]
        self._txt = '\n'.join(['{} {}'.format(self.comment_pfx, l)
                               for l in lines])

    def modify(self, fname: str, remove=False):
        _log.info('file={}'.format(fname))
        # move input file to <name>.orig
        wfname = fname + '.orig'
        shutil.move(fname, wfname)
        # re-open input filename as the output file
        f = open(wfname, 'r')
        out = open(fname, 'w')
        # re-create the file, modified
        state = 'head'
        if remove:
            for line in f:
                if state == 'head':
                    if line.strip().startswith(self.comment_minsep):
                        state = 'copyright'
                        continue
                    else:
                        out.write(line)
                elif state == 'copyright':
                    if line.strip().startswith(self.comment_minsep):
                        state = 'code'
                else:
                    out.write(line)
        else:
            lineno = 0
            ex = re.compile(r'^[ \t\f]*#'
                             '(.*?coding[:=][ \t]*[-_.a-zA-Z0-9]+|'
                             '!/.*)')
            def write_copyright():
                out.write('{}\n'.format(self.comment_sep))
                out.write(self._txt)
                out.write('\n{}\n'.format(self.comment_sep))
            for line in f:
                lineno += 1
                sline = line.strip()
                if state == 'head':
                    if sline.startswith(self.comment_minsep):
                        state = 'copyright'  # skip past this
                    elif lineno < 3 and ex.match(sline):
                        out.write(line)
                    else:
                        state = 'text'
                        write_copyright()
                        out.write(line)
                elif state == 'copyright':
                    if sline.startswith(self.comment_minsep):
                        state = 'text'
                        write_copyright()
                elif state == 'text':
                    out.write(line)
        # finalize the output
        out.close()
        # remove moved <name>.orig, the original input file
        os.unlink(wfname)


def main() -> int:
    p = argparse.ArgumentParser()
    p.add_argument('root', help='Root path from which to find files')
    p.add_argument('pattern', nargs='*', default=[],
                   help='UNIX glob-style pattern of files to match'
                        ' (default=*.py) Prefix with "~" to take complement')
    p.add_argument('-n', '--dry-run', action='store_true', dest='dry',
                   help='Do not modify files, just show which files would '
                        'be affected.')
    p.add_argument('-r', '--remove', action='store_true', dest='remove',
                   help='Remove any existing headers')
    p.add_argument('-v', '--verbose', action='count', dest='vb',
                   default=0, help='More verbose logging')
    args = p.parse_args()

    if args.vb > 1:
        _log.setLevel(logging.DEBUG)
    elif args.vb > 0:
        _log.setLevel(logging.INFO)
    else:
        _log.setLevel(logging.WARN)

    if len(args.pattern) == 0:
        patterns = None
    else:
        # sanity-check the input patterns
        for pat in args.pattern:
            if os.path.sep in pat:
                p.error('bad pattern "{}": must be a filename, not a path'
                        .format(pat))
        patterns = args.pattern
    finder = FileFinder(args.root, glob_pat=patterns)
    if len(finder) == 0:
        _log.warning('No files found from "{}" matching {}'
                  .format(args.root, '|'.join(patterns)))
        return 1
    if args.dry:
        print_files(finder)
    else:
        modifier = FileModifier(NOTICE_TEXT)
        modify_files(finder, modifier, remove=args.remove)

    return 0


if __name__ == '__main__':
    sys.exit(main())