setup.py

import os
from pathlib import Path

from setuptools import find_packages, setup


def parse_req_file(fname, initial=None):
    """Reads requires.txt file generated by setuptools and outputs a
    new/updated dict of extras as keys and corresponding lists of dependencies
    as values.

    The input file's contents are similar to a `ConfigParser` file, e.g.
    pkg_1
    pkg_2
    pkg_3

    [extras1]
    pkg_4
    pkg_5

    [extras2]
    pkg_6
    pkg_7
    """
    reqs = {} if initial is None else initial
    cline = None
    with open(fname, "r") as f:
        for line in f.readlines():
            line = line[:-1].strip()
            if len(line) == 0:
                continue
            if line[0] == "[":
                # Add new key for current extras (if missing in dict)
                cline = line[1:-1].strip()
                if cline not in reqs:
                    reqs[cline] = []
            else:
                # Only keep dependencies from extras
                if cline is not None:
                    reqs[cline].append(line)
    return reqs


def get_version(fname):
    """Reads PKG-INFO file generated by setuptools and extracts the Version
    number."""
    res = "UNK"
    with open(fname, "r") as f:
        for line in f.readlines():
            line = line[:-1]
            if line.startswith("Version:"):
                res = line.replace("Version:", "").strip()
                break
    if res in ["UNK", ""]:
        raise ValueError(f"Missing Version number in {fname}")
    return res


def read_requirements(filename: str):
    with open(filename) as requirements_file:
        import re

        def fix_url_dependencies(req: str) -> str:
            """Pip and setuptools disagree about how URL dependencies should be handled."""
            m = re.match(
                r"^(git\+)?(https|ssh)://(git@)?github\.com/([\w-]+)/(?P<name>[\w-]+)\.git",
                req,
            )
            if m is None:
                return req
            else:
                return f"{m.group('name')} @ {req}"

        requirements = []
        for line in requirements_file:
            line = line.strip()
            if line.startswith("#") or line.startswith("-e") or len(line) <= 0:
                continue
            requirements.append(fix_url_dependencies(line))
    return requirements


def _do_setup():
    base_dir = os.path.abspath(os.path.dirname(Path(__file__)))

    if not os.path.exists(
        os.path.join(base_dir, "codenav.egg-info/dependency_links.txt")
    ):
        # Build mode for sdist
        os.chdir(base_dir)

        version_path = os.path.abspath(".VERSION")
        print(version_path)
        with open(version_path, "r") as f:
            __version__ = f.readline().strip()

        # Extra dependencies for development (actually unnecessary)
        extras = {}
        if os.path.exists("dev_requirements.txt"):
            extras = {
                "dev": [
                    l.strip()
                    for l in open("dev_requirements.txt", "r").readlines()
                    if l.strip() != ""
                ]
            }
    else:
        # Install mode from sdist
        __version__ = get_version(os.path.join(base_dir, "codenav.egg-info/PKG-INFO"))
        extras = parse_req_file(os.path.join(base_dir, "codenav.egg-info/requires.txt"))

    setup(
        name="codenav",
        version=__version__,
        description=(
            "CodeNav is a LLM-powered agent that can answer queries about code."
        ),
        long_description=open("README.md").read(),
        long_description_content_type="text/markdown",
        classifiers=[
            "Intended Audience :: Science/Research",
            "Development Status :: 3 - Alpha",
            "License :: OSI Approved :: Apache Software License",
            "Programming Language :: Python :: 3",
            "Topic :: Scientific/Engineering :: Artificial Intelligence",
        ],
        keywords=["code understanding", "LLM", "large language models"],
        url="https://github.com/allenai/codenav",
        author="Allen Institute for Artificial Intelligence",
        author_email="lucaw@allenai.org",
        packages=find_packages(
            include=["codenav", "codenav.*"],
            exclude=["*.tests", "*.tests.*", "tests.*", "tests"],
        ),
        license="Apache 2.0",
        package_data={
            "codenav": ["prompts/default/*"],
        },
        install_requires=read_requirements("requirements.txt"),
        python_requires=">=3.9",
        entry_points={"console_scripts": ["codenav=codenav.codenav_run:main"]},
        # scripts=["codenav/codenav_run.py"],
        # setup_requires=["pytest-runner"],
        # tests_require=["pytest", "pytest-cov"],
        extras_require=extras,
    )


if __name__ == "__main__":
    _do_setup()