def find()

in robogym/utils/env_utils.py [0:0]

• 57 lines of code
• 20 McCabe index (conditional complexity)

    def find(cls, pattern, fun_name=None):
        if pattern.endswith("py") and os.path.exists(pattern):
            if fun_name is None:
                fun_name = "make_env"

            print(f"Loading environment from {pattern}::{fun_name}")
            module = run_path(pattern)
            make_env = module[fun_name]
            return make_env
        elif pattern.endswith((".jsonnet", ".libsonnet", ".json")) and os.path.exists(
            pattern
        ):

            def resolve_fun_name(data, fun_name):
                for elem in fun_name.split("."):
                    data = data[elem]
                return data

            print(f"Loading environment from {pattern}::{fun_name}")
            if pattern.endswith(".json"):
                with open(pattern, "r") as f:
                    data = json.load(f)
            else:
                data = json.loads(_jsonnet.evaluate_file(pattern))
            if fun_name is not None:
                resolved_data = resolve_fun_name(data, fun_name)
            else:
                # Auto-detect a working fun_name
                candidates = ["make_env", "machine_pools.evaluator.args.make_env"]
                resolved_data = None
                for candidate in candidates:
                    try:
                        resolved_data = resolve_fun_name(data, candidate)
                        break
                    except KeyError:
                        pass
                assert (
                    resolved_data is not None
                ), "could not auto-detect a function name; please provide it (e.g. via `::machine_pools.bch-b1.args.make_env`)"

            make_env = get_function(resolved_data)
            return make_env
        else:
            matching = [
                m for p in cls.ENV_PATTERNS for m in glob.glob(p, recursive=True)
            ]
            matching = [match for match in matching if match.find(pattern) > -1]
            matching = [
                match
                for match in matching
                if not os.path.basename(match).startswith("test_")
            ]
            assert len(matching) < 2, "Found multiple environments matching %s" % str(
                matching
            )
            if len(matching) == 1:
                matching = matching[0]
                if matching.endswith(".py") and fun_name is not None:
                    matching += "::" + fun_name
                return cls.find(matching)
            else:
                assert None