Fixes from upstream: (#30)

- fix an issue when resolving multiple mixed environment in parallel
  - allow the pypi-indices section in environment.yml
  - improved parsing of versions in environment.yml

metaflow_extensions/netflix_ext/cmd/environment/environment_cmd.py

@@ -55,7 +55,14 @@
 
 
 REQ_SPLIT_LINE = re.compile(r"([^~<=>]*)([~<=>]+.*)?")
-YML_SPLIT_LINE = re.compile(r"(<=|>=|=>|=<|~=|==|<|>|=)")
+
+# Allows things like:
+# pkg = <= version
+# pkg <= version
+# pkg = version
+# pkg = ==version or pkg = =version
+# In other words, the = is optional but possible
+YML_SPLIT_LINE = re.compile(r"(?:=\s)?(<=|>=|~=|==|<|>|=)")
 
 
 class CommandObj:
@@ -152,6 +159,13 @@ def cli(ctx):
     help="Root path for Conda cached information. If not set, "
     "looks for METAFLOW_CONDA_S3ROOT (for S3)",
 )
+@click.option(
+    "--quiet-file-output",
+    default=None,
+    hidden=True,
+    type=str,
+    help="Output the quiet output to this file; used for scripting",
+)
 @click.pass_context
 def environment(
     ctx: Any,
@@ -160,6 +174,7 @@ def environment(
     datastore: str,
     environment: str,
     conda_root: Optional[str],
+    quiet_file_output: Optional[str],
 ):
     if quiet:
         echo = echo_dev_null
@@ -171,6 +186,7 @@ def environment(
     obj.echo = echo
     obj.echo_always = echo_always
     obj.datastore_type = datastore
+    obj.quiet_file_output = quiet_file_output
 
     if conda_root:
         if obj.datastore_type == "s3":
@@ -453,6 +469,10 @@ def create(
 
     if obj.quiet:
         obj.echo_always(python_bin)
+        if obj.quiet_file_output:
+            with open(obj.quiet_file_output, "w") as f:
+                f.write(python_bin)
+                f.write("\n")
     else:
         obj.echo(
             "Created environment '%s' locally, activate with `%s activate %s`"
@@ -724,6 +744,12 @@ def resolve(
         if obj.quiet:
             obj.echo_always(env_id.arch)
             obj.echo_always(env.quiet_print(existing_envs.get(env.env_id)))
+            if obj.quiet_file_output:
+                with open(obj.quiet_file_output, "w") as f:
+                    f.write(env_id.arch)
+                    f.write("\n")
+                    f.write(env.quiet_print(existing_envs.get(env.env_id)))
+                    f.write("\n")
         else:
             obj.echo("### Environment for architecture %s" % env_id.arch)
             obj.echo(env.pretty_print(existing_envs.get(env.env_id)))
@@ -809,6 +835,10 @@ def show(obj, local_only: bool, arch: str, pathspec: bool, envs: Tuple[str]):
             if obj.quiet:
                 obj.echo_always("%s@%s" % (env_name, arch))
                 obj.echo_always("NOT_FOUND")
+                if obj.quiet_file_output:
+                    with open(obj.quiet_file_output, "w") as f:
+                        f.write("%s@%s\n" % (env_name, arch))
+                        f.write("NOT_FOUND\n")
             else:
                 obj.echo(
                     "### Environment for '%s' (arch '%s') was not found"
@@ -825,6 +855,17 @@ def show(obj, local_only: bool, arch: str, pathspec: bool, envs: Tuple[str]):
                         )
                     )
                 )
+                if obj.quiet_file_output:
+                    with open(obj.quiet_file_output, "w") as f:
+                        f.write("%s@%s\n" % (env_name, arch))
+                        f.write(
+                            resolved_env.quiet_print(
+                                all_envs.get(resolved_env.env_id.req_id, {}).get(
+                                    resolved_env.env_id.full_id
+                                )
+                            )
+                        )
+                        f.write("\n")
             else:
                 obj.echo("### Environment for '%s' (arch '%s'):" % (env_name, arch))
                 obj.echo(
@@ -909,6 +950,10 @@ def get(obj, default: bool, arch: Optional[str], pathspec: bool, source_env: str
         existing_envs = []
     if obj.quiet:
         obj.echo_always(env.quiet_print(existing_envs))
+        if obj.quiet_file_output:
+            with open(obj.quiet_file_output, "w") as f:
+                f.write(env.quiet_print(existing_envs))
+                f.write("\n")
     else:
         obj.echo(env.pretty_print(existing_envs))
     cast(Conda, obj.conda).write_out_environments()
@@ -1007,11 +1052,15 @@ def _parse_yml_file(
                     mode = "sources"
                 elif line == "dependencies:":
                     mode = "deps"
+                elif line == "pypi-indices:":
+                    mode = "pypi_sources"
                 else:
                     mode = "ignore"
-            elif mode == "sources":
+            elif mode == "sources" or mode == "pypi_sources":
                 line = line.lstrip(" -").rstrip()
-                sources.setdefault("conda", []).append(line)
+                sources.setdefault("conda" if mode == "sources" else "pypi", []).append(
+                    line
+                )
             elif mode == "deps" or mode == "pypi_deps":
                 line = line.lstrip(" -").rstrip()
                 if line == "pip:":

metaflow_extensions/netflix_ext/plugins/conda/conda.py

@@ -317,6 +317,7 @@ def call_binary(
         args: List[str],
         binary: str,
         addl_env: Optional[Mapping[str, str]] = None,
+        cwd: Optional[str] = None,
         pretty_print_exception: bool = True,
     ) -> bytes:
         if binary in _CONDA_DEP_RESOLVERS:
@@ -333,6 +334,7 @@ def call_binary(
                 [binary] + args,
                 stderr=subprocess.PIPE,
                 env=dict(os.environ, **addl_env),
+                cwd=cwd,
             ).strip()
         except subprocess.CalledProcessError as e:
             print(

metaflow_extensions/netflix_ext/plugins/conda/resolvers/conda_lock_resolver.py

@@ -243,20 +243,19 @@ def resolve(
                         virtual_yml.writelines(lines)
                     args.extend(["--virtual-package-spec", "virtual_yml.spec"])
 
-                with WithDir(conda_lock_dir):
-                    debug.conda_exec("Build directory: %s" % conda_lock_dir)
-                    # conda-lock will only consider a `pyproject.toml` as a TOML file which
-                    # is somewhat annoying.
-                    with open(
-                        "pyproject.toml", mode="w", encoding="ascii"
-                    ) as input_toml:
-                        input_toml.writelines(toml_lines)
-                        debug.conda_exec(
-                            "TOML configuration:\n%s" % "".join(toml_lines)
-                        )
-                    self._conda.call_binary(
-                        args, binary="conda-lock", addl_env=addl_env
-                    )
+                debug.conda_exec("Build directory: %s" % conda_lock_dir)
+                # conda-lock will only consider a `pyproject.toml` as a TOML file which
+                # is somewhat annoying.
+                with open(
+                    os.path.join(conda_lock_dir, "pyproject.toml"),
+                    mode="w",
+                    encoding="ascii",
+                ) as input_toml:
+                    input_toml.writelines(toml_lines)
+                    debug.conda_exec("TOML configuration:\n%s" % "".join(toml_lines))
+                self._conda.call_binary(
+                    args, binary="conda-lock", addl_env=addl_env, cwd=conda_lock_dir
+                )
             # At this point, we need to read the explicit dependencies in the file created
             emit = False
             packages = []  # type: List[PackageSpecification]

metaflow_extensions/netflix_ext/plugins/conda/utils.py

@@ -661,6 +661,8 @@ def change_pypi_package_version(
 
 
 class WithDir:
+    # WARNING: os.chdir is not compatible with thread processing so do not use in
+    # a context where multiple threads can exist.
     def __init__(self, new_dir: str):
         self._current_dir = os.getcwd()
         self._new_dir = new_dir