Function import_mod

tools/Polygraphy/polygraphy/mod/importer.py:117–201

()

        _all_external_lazy_imports.add(name)

    def import_mod():
        from polygraphy import config
        from polygraphy.logger import G_LOGGER, LogMode

        def install_mod(install_name, install_version, raise_error=True):
            modname = install_name.split(".")[0]
            pkg = pkg_name if pkg_name is not None else _PKG_NAME_FROM_MODULE.get(modname, modname)
            extra_flags = install_flags if install_flags is not None else _EXTRA_FLAGS_FOR_MODULE.get(modname, [])

            def fail():
                log_func = G_LOGGER.critical if raise_error else G_LOGGER.warning
                log_func(f"Could not automatically install required module: {pkg}. Please install it manually.")

            if config.ASK_BEFORE_INSTALL:
                res = None
                while res not in ["y", "n"]:
                    res = input(f"Automatically install '{pkg}' (version: {install_version or 'any'}) ([Y]/n)? ")
                    res = res.strip()[:1].lower() or "y"

                if res == "n":
                    fail()

            if install_version == LATEST_VERSION:
                extra_flags.append("--upgrade")
            elif install_version is not None:
                pkg += install_version

            cmd = config.INSTALL_CMD + [pkg] + extra_flags
            G_LOGGER.info(f"Running installation command: {' '.join(cmd)}")
            status = sp.run(cmd, stdout=sp.PIPE, stderr=sp.PIPE)
            if status.returncode != 0:
                G_LOGGER.error(f"Error during installation:\n{constants.TAB}{status.stderr.decode()}")
                fail()

            mod = importlib.import_module(install_name)
            return mod

        mod = None
        try:
            mod = importlib.import_module(name)
        except ImportError as err:
            if config.AUTOINSTALL_DEPS:
                for install_name, install_version in all_required_mods:
                    G_LOGGER.info(
                        f"Module: '{install_name}' is required, but not installed. Attempting to install now."
                    )
                    mod = install_mod(install_name, install_version)
            else:
                G_LOGGER.critical(
                    f"Module: '{name}' is required but could not be imported.\nNote: Error was: {err}\n"
                    f"You can set POLYGRAPHY_AUTOINSTALL_DEPS=1 in your environment variables to allow "
                    f"Polygraphy to automatically install missing modules.\n"
                )

        # Auto-upgrade if necessary
        for install_name, install_version in all_required_mods:
            installed_mod = importlib.import_module(install_name)
            if (