Merge pull request #302 from nickjalbert/nj_isort

agentos/__init__.py

@@ -1,6 +1,9 @@
 """The ``agentos`` module provides an API for building learning agents."""
 
-from agentos.version import VERSION as __version__  # noqa: F401
+from agentos.agent_run import AgentRun
+from agentos.argument_set import ArgumentSet
+from agentos.component import Component
+from agentos.component_run import ComponentRun, active_component_run
 from agentos.core import (
     Agent,
     Dataset,
@@ -10,23 +13,19 @@
     Runnable,
     Trainer,
 )
-
 from agentos.registry import Registry
-from agentos.component import Component
-from agentos.component_run import ComponentRun, active_component_run
-from agentos.repo import Repo, LocalRepo, GitHubRepo
-from agentos.argument_set import ArgumentSet
+from agentos.repo import GitHubRepo, LocalRepo, Repo
 from agentos.run import Run
 from agentos.run_command import RunCommand
 from agentos.specs import (
+    ArgumentSetSpec,
     ComponentSpec,
     RepoSpec,
-    ArgumentSetSpec,
     RunSpec,
     flatten_spec,
     unflatten_spec,
 )
-from agentos.agent_run import AgentRun
+from agentos.version import VERSION as __version__  # noqa: F401
 
 __all__ = [
     "Agent",

agentos/agent_run.py

@@ -1,9 +1,11 @@
 import statistics
-from typing import Optional
 from collections import namedtuple
-from agentos.run import Run
-from mlflow.utils.mlflow_tags import MLFLOW_PARENT_RUN_ID, MLFLOW_RUN_NAME
+from typing import Optional
+
 from mlflow.entities import RunStatus
+from mlflow.utils.mlflow_tags import MLFLOW_PARENT_RUN_ID, MLFLOW_RUN_NAME
+
+from agentos.run import Run
 
 _EPISODE_KEY = "episode_count"
 _STEP_KEY = "step_count"

agentos/argument_set.py

@@ -1,8 +1,10 @@
 import copy
-import yaml
 import json
 from hashlib import sha1
-from typing import TypeVar, Mapping, Dict
+from typing import Dict, Mapping, TypeVar
+
+import yaml
+
 from agentos.specs import ArgumentSetSpec
 
 # Use Python generics (https://mypy.readthedocs.io/en/stable/generics.html)

agentos/cli.py

@@ -4,14 +4,15 @@
 """
 import os
 import sys
-import yaml
-import click
 from datetime import datetime
 from pathlib import Path
-from agentos import Component
-from agentos import ArgumentSet
-from agentos.run import Run
+
+import click
+import yaml
+
+from agentos import ArgumentSet, Component
 from agentos.registry import Registry
+from agentos.run import Run
 from agentos.virtual_env import VirtualEnv
 
 

agentos/component.py

@@ -1,26 +1,25 @@
+import importlib
+import logging
 import sys
 import uuid
-import logging
-import importlib
 from hashlib import sha1
 from pathlib import Path
+from typing import Any, Dict, Sequence, Type, TypeVar, Union
+
 from dill.source import getsource as dill_getsource
-from typing import Union, TypeVar, Dict, Type, Any, Sequence
 from rich import print as rich_print
 from rich.tree import Tree
-from agentos.run_command import RunCommand
+
+from agentos.argument_set import ArgumentSet
 from agentos.component_run import ComponentRun
+from agentos.exceptions import RegistryException
 from agentos.identifiers import ComponentIdentifier
+from agentos.registry import InMemoryRegistry, Registry
+from agentos.repo import GitHubRepo, LocalRepo, Repo, RepoType
+from agentos.run_command import RunCommand
 from agentos.specs import ComponentSpec, ComponentSpecKeys, unflatten_spec
-from agentos.registry import (
-    Registry,
-    InMemoryRegistry,
-)
-from agentos.exceptions import RegistryException
-from agentos.repo import RepoType, Repo, LocalRepo, GitHubRepo
-from agentos.argument_set import ArgumentSet
-from agentos.virtual_env import VirtualEnv, NoOpVirtualEnv
 from agentos.utils import parse_github_web_ui_url
+from agentos.virtual_env import NoOpVirtualEnv, VirtualEnv
 
 logger = logging.getLogger(__name__)
 

agentos/component_run.py

@@ -1,7 +1,9 @@
 import tempfile
-from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME
 from pathlib import Path
-from typing import Any, Optional, Mapping
+from typing import Any, Mapping, Optional
+
+from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME
+
 from agentos.exceptions import PythonComponentSystemException
 from agentos.identifiers import RunIdentifier
 from agentos.registry import Registry

agentos/core.py

@@ -2,7 +2,8 @@
 import time
 from collections import namedtuple
 from threading import Thread
-from typing import Sequence, Optional
+from typing import Optional, Sequence
+
 from agentos.agent_run import AgentRun
 
 

agentos/identifiers.py

@@ -5,6 +5,7 @@ class ComponentIdentifier(str):
     allows referring to Components both as [name] and [name]==[version] (e.g.,
     in registries or the command line).
     """
+
     def __new__(
         cls,
         identifier_or_name: str,

agentos/registry.py

@@ -1,40 +1,42 @@
 import abc
+import json
 import logging
 import os
-import yaml
-import json
 import pprint
 import shutil
 import tarfile
 import tempfile
-import requests
 from pathlib import Path
-from typing import Dict, Sequence, Union, Optional, TYPE_CHECKING
+from typing import TYPE_CHECKING, Dict, Optional, Sequence, Union
+
+import requests
+import yaml
 from dotenv import load_dotenv
+
 from agentos.identifiers import (
     ComponentIdentifier,
-    RunIdentifier,
     RepoIdentifier,
     RunCommandIdentifier,
+    RunIdentifier,
 )
 from agentos.specs import (
-    flatten_spec,
-    unflatten_spec,
-    is_flat,
-    json_encode_flat_spec_field,
-    RepoSpec,
     ComponentSpec,
     NestedComponentSpec,
-    RunSpec,
+    RepoSpec,
     RunCommandSpec,
+    RunSpec,
+    flatten_spec,
+    is_flat,
+    json_encode_flat_spec_field,
+    unflatten_spec,
 )
 
 logger = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
     from agentos.component import Component
-    from agentos.run import Run
     from agentos.repo import Repo
+    from agentos.run import Run
 
 # add USE_LOCAL_SERVER=True to .env to talk to local server
 load_dotenv()
@@ -194,8 +196,7 @@ def get_component_spec(
                 return {}
         if len(components) > 1:
             versions = [
-                ComponentIdentifier(c_id).version
-                for c_id in components.keys()
+                ComponentIdentifier(c_id).version for c_id in components.keys()
             ]
             version_str = "\n - ".join(versions)
             raise LookupError(
@@ -233,9 +234,7 @@ def get_specs_transitively_by_id(
             repo_spec = self.get_repo_spec(repo_id, flatten=flatten)
             repo_specs[repo_id] = repo_spec
             for d_id in inner_spec.get("dependencies", {}).values():
-                component_identifiers.append(
-                    ComponentIdentifier(d_id)
-                )
+                component_identifiers.append(ComponentIdentifier(d_id))
         return list(component_specs.values()), list(repo_specs.values())
 
     def has_component_by_id(self, identifier: ComponentIdentifier) -> bool:

agentos/repo.py

@@ -1,25 +1,25 @@
-import os
-import sys
 import abc
 import logging
+import os
+import sys
 import uuid
 from enum import Enum
-from typing import TypeVar, Dict, Tuple, Union
 from pathlib import Path
+from typing import Dict, Tuple, TypeVar, Union
+
 from dulwich import porcelain
-from dulwich.repo import Repo as PorcelainRepo
-from dulwich.objectspec import parse_ref
-from dulwich.objectspec import parse_commit
 from dulwich.errors import NotGitRepository
+from dulwich.objectspec import parse_commit, parse_ref
+from dulwich.repo import Repo as PorcelainRepo
 
 from agentos.exceptions import (
     BadGitStateException,
     PythonComponentSystemException,
 )
-from agentos.utils import AOS_GLOBAL_REPOS_DIR
 from agentos.identifiers import ComponentIdentifier, RepoIdentifier
-from agentos.specs import RepoSpec, NestedRepoSpec, RepoSpecKeys, flatten_spec
-from agentos.registry import Registry, InMemoryRegistry
+from agentos.registry import InMemoryRegistry, Registry
+from agentos.specs import NestedRepoSpec, RepoSpec, RepoSpecKeys, flatten_spec
+from agentos.utils import AOS_GLOBAL_REPOS_DIR
 
 logger = logging.getLogger(__name__)
 

agentos/run.py

@@ -3,14 +3,15 @@
 from pathlib import Path
 from typing import Sequence
 from urllib.parse import urlparse
+
 from mlflow.entities import RunStatus
 from mlflow.exceptions import MlflowException
 from mlflow.tracking import MlflowClient
 from mlflow.tracking.context import registry as context_registry
+
 from agentos.identifiers import RunIdentifier
 from agentos.registry import Registry
-from agentos.specs import RunSpec
-from agentos.specs import unflatten_spec
+from agentos.specs import RunSpec, unflatten_spec
 
 
 class Run:

agentos/run_command.py

@@ -1,14 +1,15 @@
 from hashlib import sha1
 from typing import TYPE_CHECKING
+
+from agentos.identifiers import RunCommandIdentifier, RunIdentifier
 from agentos.registry import Registry
-from agentos.specs import RunCommandSpec, RunCommandSpecKeys, unflatten_spec
-from agentos.identifiers import RunIdentifier, RunCommandIdentifier
 from agentos.run import Run
+from agentos.specs import RunCommandSpec, RunCommandSpecKeys, unflatten_spec
 
 # Avoids circular imports
 if TYPE_CHECKING:
-    from agentos.component import Component
     from agentos.argument_set import ArgumentSet
+    from agentos.component import Component
 
 
 class RunCommand:
@@ -115,8 +116,8 @@ def from_spec(
             spec_identifier = key
             inner_spec = value
         component_id = inner_spec[RunCommandSpecKeys.COMPONENT_ID]
-        from agentos.component import Component
         from agentos.argument_set import ArgumentSet
+        from agentos.component import Component
 
         component = Component.from_registry(registry, component_id)
         arg_set = ArgumentSet.from_spec(

agentos/specs.py

@@ -19,7 +19,8 @@
 
 import copy
 import json
-from typing import Mapping, Union, Any
+from typing import Any, Mapping, Union
+
 from agentos.identifiers import ComponentIdentifier
 
 FlatSpec = Mapping[str, str]

agentos/utils.py

@@ -1,8 +1,8 @@
 import pprint
-import yaml
 from pathlib import Path
 from typing import Dict, Optional
 
+import yaml
 
 AOS_GLOBAL_CONFIG_DIR = Path.home() / ".agentos"
 AOS_GLOBAL_CACHE_DIR = AOS_GLOBAL_CONFIG_DIR / "cache"

agentos/virtual_env.py

@@ -1,13 +1,14 @@
+import hashlib
 import os
-import sys
-import yaml
 import shutil
-import hashlib
-import sysconfig
 import subprocess
+import sys
+import sysconfig
+from contextlib import contextmanager
 from pathlib import Path
 from typing import Sequence
-from contextlib import contextmanager
+
+import yaml
 
 from agentos.utils import AOS_GLOBAL_REQS_DIR
 

dev-requirements.txt

@@ -13,6 +13,7 @@ sphinx_rtd_theme==1.0.0
 click>=7.0 # when updating, also update in ../setup.py
 
 # Tests requirements
+isort==5.10.1
 black==21.12b0
 cloudpickle==1.3.0  # gym 0.17.1 in setup.py requires cloudpickle<1.4.0,>=1.2.0
 flake8==4.0.1

documentation/conf.py

@@ -13,7 +13,6 @@
 import os
 from importlib.machinery import SourceFileLoader
 
-
 version = (
     SourceFileLoader(
         "agentos.version", os.path.join("..", "agentos", "version.py")

example_agents/acme_dqn/agent.py

@@ -1,6 +1,7 @@
 import acme
-from acme.agents.tf import dqn
 import numpy as np
+from acme.agents.tf import dqn
+
 from agentos import active_component_run
 
 

example_agents/acme_dqn/environment.py

@@ -1,11 +1,10 @@
 # Thin wrapper around cartpole from OpenAI's Gym toolkit
 # This env models a cart with a pole balancing on top of it
-import agentos
 import gym
 import numpy as np
-from dm_env import specs
-from dm_env import TimeStep
-from dm_env import StepType
+from dm_env import StepType, TimeStep, specs
+
+import agentos
 
 
 class CartPole(agentos.Environment):