setup.py

import os
import platform
import subprocess
import sys
from os import path

from setuptools import Extension, find_packages, setup
from setuptools.command.develop import develop
from setuptools.command.egg_info import egg_info
from setuptools.command.install import install

try:
    from setuptools import setup
except ImportError:
    from distutils.core import setup

if sys.version_info < (3, 10, 0):
    raise OSError(f"Marie requires Python >=3.10, but yours is {sys.version}")

try:
    # marie was already taken by PIP
    pkg_name = "marie-ai"
    lib_name = "marie"
    libinfo_py = path.join(lib_name, "__init__.py")
    libinfo_content = open(libinfo_py, "r", encoding="utf8").readlines()
    version_line = [l.strip() for l in libinfo_content if l.startswith("__version__")][
        0
    ]
    exec(version_line)  # gives __version__
except FileNotFoundError:
    __version__ = '0.0.0'

try:
    with open('README.md', encoding='utf-8') as fp:
        _long_description = fp.read()
except FileNotFoundError:
    _long_description = ""


def register_ac():
    import os
    import re
    from pathlib import Path

    home = str(Path.home())
    resource_path = "marie/resources/completions/marie.%s"
    regex = r"#\sMARIE_CLI_BEGIN(.*)#\sMARIE_CLI_END"
    _check = {"zsh": ".zshrc", "bash": ".bashrc", "fish": ".fish"}

    def add_ac(k, v):
        v_fp = os.path.join(home, v)
        if os.path.exists(v_fp) or os.environ.get('SHELL', '').endswith(k):
            try:
                with open(v_fp, encoding='utf-8') as fp:
                    sh_content = fp.read()
            except FileNotFoundError:
                sh_content = ''
            with open(resource_path % k, encoding='utf-8') as fr:
                if re.findall(regex, sh_content, flags=re.S):
                    _sh_content = re.sub(regex, fr.read(), sh_content, flags=re.S)
                else:
                    _sh_content = sh_content + '\n\n' + fr.read()
            if _sh_content:
                with open(v_fp, 'w', encoding='utf-8') as fp:
                    fp.write(_sh_content)

    try:
        for k, v in _check.items():
            add_ac(k, v)
    except Exception:
        pass


class PostDevelopCommand(develop):
    """Post-installation for development mode."""

    def run(self):
        develop.run(self)
        register_ac()


class PostInstallCommand(install):
    """Post-installation for installation mode."""

    def run(self):
        install.run(self)
        register_ac()


class PostEggInfoCommand(egg_info):
    """Post-installation for egg info mode."""

    def run(self):
        egg_info.run(self)
        register_ac()


def get_extra_requires(path, add_all=True):
    import re
    from collections import defaultdict

    try:
        with open(path, encoding='utf-8') as fp:
            extra_deps = defaultdict(set)
            for k in fp:
                if k.strip() and not k.startswith("#"):
                    tags = set()
                    if ":" in k:
                        rpos = k.rindex(":")
                        k, v = (k[0:rpos], k[rpos + 1 : len(k)])
                        print(f"++ {k}  ===== {v}")
                        # k, v = k.split(":")
                        tags.update(vv.strip() for vv in v.split(","))
                    tags.add(re.split("[<=>]", k)[0])
                    for t in tags:
                        extra_deps[t].add(k)

            # add tag `all` at the end
            if add_all:
                extra_deps['all'] = set(vv for v in extra_deps.values() for vv in v)

        return extra_deps
    except FileNotFoundError:
        return {}


all_deps = get_extra_requires('extra-requirements.txt')

core_deps = all_deps['core']
perf_deps = all_deps['perf'].union(core_deps)
standard_deps = all_deps['standard'].union(core_deps).union(perf_deps)

# uvloop is not supported on windows
perf_deps = {
    i + ";platform_system!='Windows'" if i.startswith('uvloop') else i
    for i in perf_deps
}
standard_deps = {
    i + ";platform_system!='Windows'" if i.startswith('uvloop') else i
    for i in standard_deps
}
for k in ['all', 'devel', 'cicd']:
    all_deps[k] = {
        i + ";platform_system!='Windows'" if i.startswith('uvloop') else i
        for i in all_deps[k]
    }

# by default, final deps is the standard deps, unless specified by env otherwise
final_deps = standard_deps

# Use env var to enable a minimum installation of Marie
# MARIE_PIP_INSTALL_CORE=1 pip install marie
# MARIE_PIP_INSTALL_PERF=1 pip install marie
if os.environ.get("MARIE_PIP_INSTALL_CORE"):
    final_deps = core_deps
elif os.environ.get("MARIE_PIP_INSTALL_PERF"):
    final_deps = perf_deps

if sys.version_info.major == 3 and sys.version_info.minor >= 11:
    for dep in list(final_deps):
        if dep.startswith('grpcio'):
            final_deps.remove(dep)
    final_deps.add('grpcio>=1.49.0')
    final_deps.add('grpcio-health-checking>=1.49.0')
    final_deps.add('grpcio-reflection>=1.49.0')


extra_golang_kw = {}

ret_code = -1

try:
    ret_code = subprocess.run(['go', 'version']).returncode
except Exception:
    pass

is_mac_os = platform.system() == 'Darwin'
is_windows_os = platform.system() == 'Windows'
is_37 = sys.version_info.major == 3 and sys.version_info.minor == 7

if ret_code == 0 and not is_windows_os and (not is_mac_os or not is_37):
    extra_golang_kw = {
        'build_golang': {'root': 'jraft', 'strip': False},
        'ext_modules': [
            Extension(
                'jraft',
                ['marie/serve/consensus/run.go'],
                py_limited_api=True,
                define_macros=[('Py_LIMITED_API', None)],
            )
        ],
        'setup_requires': ['setuptools-golang'],
    }

setup(
    name=pkg_name,
    packages=find_packages(),
    version=__version__,
    include_package_data=True,
    description="Python library to Integrate AI-powered features into your applications",
    author="Marie AI",
    author_email="hello@marieai.co",
    license="Apache 2.0",
    url="https://github.com/marieai/marie-ai/",
    download_url="https://github.com/marieai/marie-ai/tags",
    long_description=_long_description,
    long_description_content_type="text/markdown",
    zip_safe=False,
    install_requires=list(final_deps),
    extras_require=all_deps,
    entry_points={
        "console_scripts": [
            "marie=marie_cli:main",
        ],
    },
    cmdclass={
        'develop': PostDevelopCommand,
        'install': PostInstallCommand,
        'egg_info': PostEggInfoCommand,
    },
    classifiers=[
        "Development Status :: 5 - Production/Stable",
        "Intended Audience :: Developers",
        "Intended Audience :: Education",
        "Intended Audience :: Science/Research",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: 3.11",
        "Programming Language :: Unix Shell",
        "Environment :: Console",
        "License :: OSI Approved :: Apache Software License",
        "Operating System :: OS Independent",
        "Topic :: Database :: Database Engines/Servers",
        "Topic :: Scientific/Engineering :: Artificial Intelligence",
        "Topic :: Internet :: WWW/HTTP :: Indexing/Search",
        "Topic :: Scientific/Engineering :: Image Recognition",
        "Topic :: Multimedia :: Video",
        "Topic :: Scientific/Engineering",
        "Topic :: Scientific/Engineering :: Mathematics",
        "Topic :: Software Development",
        "Topic :: Software Development :: Libraries",
        "Topic :: Software Development :: Libraries :: Python Modules",
    ],
    project_urls={
        "Documentation": "https://docs.marieai.co",
        "Source": "https://github.com/marieai/marie-ai.git",
        "Tracker": "https://github.com/marieai/marie-ai/issues",
    },
    keywords=(
        "marie-ai ocr icr index elastic neural-network encoding "
        "embedding serving docker container image video audio deep-learning mlops"
    ),
)