langgrep

#! /usr/bin/env python3

# langgrep: grep for a pattern but only in files written in the
# specified language (as specified by the shebang line).
#
# Copyright 2009, 2019 by Akkana Peck.
# Please share, modify and enjoy under the terms of the GPL v2
# or, at your option, any later GPL version.
#
# Bugs/To Do:
# 1. it isn't smart about parsing the grep flags.
# Anything beginning with a - will be considered a flag
# and passed on to grep; the first argument not starting with -
# is taken to be the search pattern, and everything after that
# is the file list.
#
# 2. If you have files in two directories in the search path,
# e.g. ~/bin/word2html -> ~/src/scripts/word2html.py,
# you'll see results twice. It should check for symlinks and
# omit anything that links to another directory in the search path.

import string, os, sys
import subprocess
import shlex
from pathlib import Path


# Use XDG for the config and cache directories if it's available:
try:
    import xdg.BaseDirectory
except:
    pass


def Usage():
    print("langgrep [-f] lang [grepflags] pattern [files]")
    print("     e.g. langgrep python -w find")
    print("-f: print full path of files")
    print()
    print("If no files are specified, will search in ~/bin")
    print("plus any directory specified in ~/.config/langgrep/$LANGUAGE.conf")
    sys.exit(0)


def extra_dirs_by_lang(lang):
    """See if there's a lang.conf file that lists extra files by language.
       It should have absolute paths to directories, one per line.
       Return a list of extra dirs.
    """
    #
    # Read the config file
    #
    if 'XDG_CONFIG_HOME' in os.environ:
        confighome = os.environ['XDG_CONFIG_HOME']
    elif 'xdg.BaseDirectory' in sys.modules:
        confighome = xdg.BaseDirectory.xdg_config_home
    else:
        confighome = os.path.join(os.environ['HOME'], '.config')

    confdir = os.path.join(confighome, 'langgrep')
    configfile = os.path.join(confdir, '%s.conf' % lang)
    # print("configfile:", configfile)

    extradirs = []
    try:
        with open(configfile) as cf:
            for line in cf:
                # Allow $HOME or ~ in specifiers
                if '$HOME' in line:
                    line = line.replace('$HOME', '~')
                line = os.path.expanduser(line.strip())
                if line and not line.startswith('#') and os.path.exists(line):
                    extradirs.append(line)
    except:
        pass

    # print(extradirs)
    return extradirs


file_endings = {
    'python' : [ '.py', 'py3' ],
    'javascript' : [ '.js' ],
    'js' : [ '.js' ],
    'php' : [ '.php' ],
    'c' : [ '.c', '.h' ],
    'c++' : [ '.cpp', '.c++', '.h' ],
    'java' : [ '.java' ],
    'ruby' : [ '.rb' ],
    'perl' : [ '.pl', 'perl' ],
    'sh' : [ '.sh' ],
    'bash' : [ '.bash' ],
    'zsh' : [ '.zsh' ],
    'csh' : [ '.csh' ],
}

languages_with_shebang = [ 'python', 'perl', 'sh', 'ruby', 'perl',
                           'bash', 'zsh', 'csh', 'tcsh' ]


def check_file_lang(filename, lang):
    """Try to guess whether a file is a given programming language.
       Use file extensions, shebangs.
    """

    if os.path.isdir(filename) or not os.path.exists(filename):
        return False

    if lang in file_endings:
        for ending in file_endings[lang]:
            if filename.endswith(ending):
                return True

    # Don't waste time looking for shebangs in every file
    # for languages where that's not relevant:
    if lang in languages_with_shebang:
        try:
            f = open(filename, 'r')
            firstline = f.readline()
            f.close()
        except IOError as e:
            print("IOError", e)
            # print "exc_info is", sys.exc_info()
            sys.exit(1)
            return False
        except UnicodeDecodeError as e:
            # print(filename, "seems to be a binary file")
            return False

        if firstline[0:2] == "#!" and lang in firstline:
            return True

    return False


# Keep track of filepaths already seen, so as not to repeat them.
filepaths = set()

def find_files_in_dir(lang, direc):
    """Generator: iterate over a set of files that are the right language
       under the given dir.
    """

    for root, dirs, fs in os.walk(direc, followlinks=True):
        for f in fs:
            filename = os.path.join(root, f)
            if check_file_lang(filename, lang):
                filepath = os.path.join(root, filename)
                if filepath not in filepaths:
                    filepaths.add(filepath)
                    yield filepath


def langgrep(lang, pattern, grepargs, fil, flags):
    arglist = ['grep', '-H']
    arglist.extend(grepargs)
    arglist.append(pattern)
    arglist.append(fil)
    proc = subprocess.Popen(arglist,
                            shell=False, stdout=subprocess.PIPE)
    pout = proc.communicate()[0]

    # Go through the output removing all but the last dir of the path.
    # Otherwise lines are so long they're confusing.
    if not pout:
        return

    # for line in pout.decode("utf-8").split('\n'):
    for line in pout.decode().split('\n'):
        # For some reason the split is giving us every other line empty
        if not line:
            continue

        # If there's a super long line, it's not part of my code
        # and won't be easy to read, so omit it:
        if len(line) > 160:
            continue

        colon = line.find(":")
        pathparts = line[:colon].split('/')
        if len(pathparts) < 2 or pathparts[-2] == "bin":
            path = pathparts[-1]
        else:
            path = '/'.join(pathparts[-2:])

        if 'fullpath' in flags and flags['fullpath']:
            # XXX This isn't really the full path, needs work
            s = "%s%s" % (path, line[colon:])
        else:
            s = "%s%s" % (os.path.basename(path), line[colon:])
        # print(s.encode('utf-8', "backslashreplace"))
        print(s)


def is_link_to_dirs(filepath, alldirs):
    """Is filepath a symbolic link to anything inside any of the
       directories in alldirs?
    """
    p = Path(filepath)

    # Is it a symlink? There must be a better way to check this
    # than statting twice.
    if p.stat().st_ino == p.lstat().st_ino:
        return False

    # It is a symlink. Does it link to any of the given directories?
    parentdir = str(p.parent)
    for d in alldirs:
        if d == parentdir:
            # print("Skipping symlink", filepath)
            return True
    return False


def parse_args():
    """Usage: langgrep lang [grepflags] pattern files
       Returns: (lang, pattern, grepargs, files)
    """
    if len(sys.argv) < 3:
        Usage()

    flags = {}

    if sys.argv[1] == '-f':
        flags['fullpath'] = True
        sys.argv = sys.argv[1:]

    elif sys.argv[1].startswith('-'):
        Usage()

    lang = sys.argv[1]

    # After the language, any flag argument plus the grep pattern
    # gets appended to args -- these will be the grep args.
    patindex = 0

    grepargs = []
    pattern = None
    files = []

    for arg in sys.argv[2:]:
        if not pattern:
            if arg.startswith('-'):
                grepargs.append(arg)
            else:
                pattern = arg
            continue

        files.append(arg)

    return (lang, pattern, grepargs, files, flags)


if __name__ == '__main__':
    try:
        lang, pattern, grepargs, files, flags = parse_args()

        if not files:
            alldirs = extra_dirs_by_lang(lang)
            # ~/bin goes last, because it's most likely to contain symlinks
            alldirs.append(os.path.join(os.getenv("HOME"), "bin"))

            for d in alldirs:
                for filepath in find_files_in_dir(lang, d):
                    if is_link_to_dirs(filepath, alldirs):
                        continue
                    langgrep(lang, pattern, grepargs, filepath, flags)

        # print("files:", '\n'.join(files))

        for fil in files:
            langgrep(lang, pattern, grepargs, fil, ('fullpath') in flags)

        # Try to catch ctrl-C and print a nicer message.
        # This doesn't work, though:
        # subprocess.call just terminates the whole process.
    except KeyboardInterrupt as e:
        print("Interrupt!")
        sys.exit(1)