search_pypi_top.py

#!/usr/bin/python3 -u
"""
Code search in the source code of PyPI top projects.

Usage::

    ./search_pypi_top.py --verbose PYPI_DIR/ REGEX -o output_file

Use --help command line option to get the command line usage.

Faster alternative but producing more false alarms:

    rg -zl "REGEX" pypi_directory/*.{zip,gz,bz2,tgz}
"""
import argparse
import datetime
import logging
import os
import re
import sys
import tarfile
import zipfile
from itertools import repeat
import multiprocessing

try:
    from termcolor import colored
except ImportError:
    print(
        "Warning: termcolor is missing, install it with: "
        "python -m pip install --user termcolor",
        file=sys.stderr,
    )
    print(file=sys.stderr, end="")

    def colored(msg, *ignored, **ignored2):
        return msg


IGNORE_CYTHON = True
# Ignore macOS directory metadata files from the initial directory listing
DS_STORE = ".DS_Store"
# Ignore file extensions known to be binary files to avoid the slow
# is_binary_string() check
IGNORED_FILE_EXTENSIONS = (
    # Programs and dynamic libraries
    "EXE", "SO", "PYD",
    # Python
    "PYC", "WHL",
    # Archives
    "BZ2", "CAB", "GZ", "ZIP", "RAR", "TAR", "TGZ", "XZ",
    # Pictures
    "AI", "BMP", "ICO", "JPG", "PNG", "PSD", "SVG", "TGA", "TIF", "WMF", "XCF",
    # Audio
    "AAC", "AIF", "ANI", "MP3", "WAV", "WMA",
    # Video
    "AVI", "MKV", "MP4",
    # Linux packages
    "DEB", "RPM",
    # Misc
    "BIN", "PDF", "MO", "DB", "ISO", "JAR", "TTF", "XLS",
    "DS_Store",
    # Text which is not Python nor C code but can contain false positive,
    # and is unlikely to be used as template to generate Python or C code
    "JS", "JS.MAP", "JSON", "RST", "HTML", "CSS",
)
IGNORED_FILE_EXTENSIONS = tuple("." + ext.lower()
                                for ext in IGNORED_FILE_EXTENSIONS)
# Check the first bytes of a file to test if it's a binary file or not
BINARY_TEST_LEN = 256

# "/* Generated by Cython 0.29.13 */"
# "/* Generated by Cython 0.20.1 on Sun Mar 16 22:58:12 2014 */"
CYTHON_REGEX = re.compile(br'^/\* Generated by Cython [0-9]+(.[0-9]+)+ ')

formatter = logging.Formatter("# %(message)s")
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logger = multiprocessing.get_logger()
logger.addHandler(handler)


def ignore_filename(args, filename):
    if args.text:
        return False
    return filename.lower().endswith(IGNORED_FILE_EXTENSIONS)


def log_ignored_file(archive_name, filename):
    logger.info(f"ignore filename: {archive_name}: {filename}")


# If all bytes of a string are in TEXTCHARS, the string looks like text
TEXTCHARS = bytearray({7,8,9,10,12,13,27} | set(range(0x20, 0x100)) - {0x7f})


def is_binary_string(header):
    # Treat UTF-16 as binary, but that's acceptable since Python and C source
    # files are not encoded as UTF-16
    return bool(header.translate(None, TEXTCHARS))


def decompress_tar(args, archive_filename, mode):
    with tarfile.open(archive_filename, mode) as tar:
        while True:
            member = tar.next()
            if member is None:
                break
            filename = member.name
            if ignore_filename(args, filename):
                log_ignored_file(archive_filename, filename)
                continue
            try:
                fp = tar.extractfile(member)
                if fp is None:
                    continue
            except KeyError:
                # graphitesend-0.10.0.tar.gz fails with:
                #  File "tarfile.py", line 2124, in extractfile
                #    return self.extractfile(self._find_link_target(tarinfo))
                #  File "tarfile.py", line 2431, in _find_link_target
                #    raise KeyError("linkname %r not found" % linkname)
                # KeyError: "linkname 'graphitesend-0.10.0/README.md' not found"
                continue

            with fp:
                yield filename, fp


def decompress_zip(args, archive_filename):
    with zipfile.ZipFile(archive_filename) as zf:
        for member in zf.filelist:
            filename = member.filename
            if ignore_filename(args, filename):
                log_ignored_file(archive_filename, filename)
                continue
            with zf.open(member) as fp:
                yield filename, fp


def decompress(args, filename):
    if filename.endswith((".tar.gz", ".tgz")):
        yield from decompress_tar(args, filename, "r:gz")
    elif filename.endswith(".tar.bz2"):
        yield from decompress_tar(args, filename, "r:bz2")
    elif filename.endswith(".zip"):
        yield from decompress_zip(args, filename)
    else:
        raise Exception(f"unsupported filename: {filename!r}")


def is_binary_file(args, fp):
    if args.text:
        return False
    data = fp.read(BINARY_TEST_LEN)
    fp.seek(0)
    if is_binary_string(data):
        return True
    return False


def grep(args, archive_filename, regex):
    for filename, fp in decompress(args, archive_filename):
        if is_binary_file(args, fp):
            logger.info(f"ignore binary file: {archive_filename}: {filename}")
            continue

        matches = []
        ignore = False
        lineno = 1
        # Split at Unix newline b'\n' byte
        for line in fp:
            if lineno == 1 and IGNORE_CYTHON and CYTHON_REGEX.match(line):
                logger.info(f"ignore Cython file: {archive_filename}: {filename}")
                ignore = True
                break
            match = regex.search(line)
            if match:
                matches.append((filename, line, match.span()))
            lineno += 1

        if matches and not ignore:
            yield from matches


def search_file(filename, index, len_filenames, args, pypi_dir, regex):
    lines = 0
    results = []
    filename = os.path.join(pypi_dir, filename)
    percent = index * 100 / len_filenames
    logger.warning(f"grep {filename} ({percent:.0f}%, {index}/{len_filenames})")

    for name, line, span in grep(args, filename, regex):
        line = line.decode('utf8', 'replace')

        # print to terminal with color
        start, end = span
        line_color = (
            line[:start] + colored(line[start:end], "red", attrs=["bold"]) + line[end:]
        )
        name_color = colored(name, "magenta")
        result = f"{filename}: {name_color}: {line_color.strip()}"
        print(result, flush=True)

        # print to file without color
        result = f"{filename}: {name}: {line.strip()}"
        results.append(result)
        lines += 1

    return lines, filename if lines >= 1 else None, results


def search_dir(args, pypi_dir, pattern):
    lines = 0
    projects = set()
    all_results = []
    regex = re.compile(pattern)
    filenames = (filename for filename in os.listdir(pypi_dir) if filename != DS_STORE)
    filenames = sorted(filenames, key=str.lower)

    with multiprocessing.Pool() as pool:
        ret = pool.starmap(
            search_file,
            zip(
                filenames,
                range(len(filenames)),
                repeat(len(filenames)),
                repeat(args),
                repeat(pypi_dir),
                repeat(regex),
            )
        )
        for new_lines, new_project, results in ret:
            lines += new_lines
            if new_project:
                projects.add(new_project)
            all_results.extend(results)
    pool.close()
    pool.join()

    return lines, projects, all_results


def parse_args():
    parser = argparse.ArgumentParser(description='Code search in the source code of PyPI top projects.')
    parser.add_argument('pypi_dir', metavar="PYPI_DIRECTORY",
                        help='PyPI local directory')
    parser.add_argument('pattern', metavar='REGEX',
                        help='Regex to search')
    parser.add_argument('-o', '--output', metavar='FILENAME',
                        help='Output filename')
    parser.add_argument('--text', action='store_true',
                        help='Process a binary file as if it were text')
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='Verbose mode (ex: log ignored files)')
    parser.add_argument('-q', '--quiet', action='store_true',
                        help="Quiet mode (ex: don't log proceed files)")
    parser.add_argument('--cython', action='store_true',
                        help="Search also in code generated by Cython")

    return parser.parse_args()


def _main():
    args = parse_args()
    output_filename = args.output
    pattern = os.fsencode(args.pattern)
    pypi_dir = args.pypi_dir
    if args.cython:
        global IGNORE_CYTHON
        IGNORE_CYTHON = False

    if args.quiet:
        level = logging.ERROR
    elif args.verbose:
        level = logging.INFO
    else:
        level = logging.WARNING
    logger.setLevel(level)

    start_time = datetime.datetime.now()
    lines, projects, results = search_dir(args, pypi_dir, pattern)

    if output_filename:
        with open(output_filename, "w", encoding="utf8") as output:
            print("\n".join(results), file=output, flush=True)

    dt = datetime.datetime.now() - start_time
    print()
    print(f"Time: {dt}")
    print(f"Found {lines} matching lines in {len(projects)} projects")
    if output_filename:
        print(f"Output written into: {output_filename}")


def main():
    try:
        _main()
    except KeyboardInterrupt:
        print()
        print("Interrupted")
        sys.exit(1)


if __name__ == "__main__":
    main()