From af740dfedb02221657c2a9779517e2bb3d693268 Mon Sep 17 00:00:00 2001
From: Eashwar Ranganathan <eashwar@eashwar.com>
Date: Sun, 4 Apr 2021 03:40:43 -0700
Subject: [PATCH] Add ability to specify rust language version

---
 pre_commit/languages/rust.py | 82 +++++++++++++++++++++++++++---------
 1 file changed, 62 insertions(+), 20 deletions(-)

diff --git a/pre_commit/languages/rust.py b/pre_commit/languages/rust.py
index 7ea3f5406..a66a3f134 100644
--- a/pre_commit/languages/rust.py
+++ b/pre_commit/languages/rust.py
@@ -18,25 +18,61 @@
 from pre_commit.util import cmd_output_b
 
 ENVIRONMENT_DIR = 'rustenv'
+RUNTIME_DIR = 'rustup'
 get_default_version = helpers.basic_get_default_version
-healthy = helpers.basic_healthy
 
 
-def get_env_patch(target_dir: str) -> PatchesT:
-    return (
-        ('PATH', (os.path.join(target_dir, 'bin'), os.pathsep, Var('PATH'))),
+def _envdir(prefix: Prefix, version: str) -> str:
+    directory = helpers.environment_dir(ENVIRONMENT_DIR, version)
+    return prefix.path(directory)
+
+
+def _version_flag(version: str) -> Sequence[str]:
+    return [] if version == C.DEFAULT else [f'+{version}']
+
+
+def _resolve_version(prefix: Prefix, version: str) -> str:
+    if version != C.DEFAULT:
+        return version
+
+    if prefix.exists('rust-toolchain'):
+        with open(prefix.path('rust-toolchain')) as f:
+            return f.readline().strip()
+
+    return version
+
+
+def get_env_patch(prefix: Prefix, version: str) -> PatchesT:
+    env_path = _envdir(prefix, version)
+    patch = (
+        ('CARGO_HOME', env_path),
+        ('PATH', (os.path.join(env_path, 'bin'), os.pathsep, Var('PATH'))),
     )
 
+    if version != C.DEFAULT:
+        return (*patch, ('RUSTUP_HOME', prefix.path(RUNTIME_DIR)))
+
+    return patch
+
 
 @contextlib.contextmanager
-def in_env(prefix: Prefix) -> Generator[None, None, None]:
-    target_dir = prefix.path(
-        helpers.environment_dir(ENVIRONMENT_DIR, C.DEFAULT),
-    )
-    with envcontext(get_env_patch(target_dir)):
+def in_env(prefix: Prefix, version: str) -> Generator[None, None, None]:
+    with envcontext(get_env_patch(prefix, version)):
         yield
 
 
+def healthy(prefix: Prefix, version: str) -> bool:
+    language_version = _resolve_version(prefix, version)
+    with in_env(prefix, language_version):
+        version_flag = _version_flag(language_version)
+        retcode, _, _ = cmd_output_b(
+            'rustc', *version_flag, '--version',
+            retcode=None,
+            cwd=prefix.prefix_dir,
+        )
+        return retcode == 0
+
+
 def _add_dependencies(
         cargo_toml_path: str,
         additional_dependencies: Set[str],
@@ -57,11 +93,6 @@ def install_environment(
         version: str,
         additional_dependencies: Sequence[str],
 ) -> None:
-    helpers.assert_version_default('rust', version)
-    directory = prefix.path(
-        helpers.environment_dir(ENVIRONMENT_DIR, C.DEFAULT),
-    )
-
     # There are two cases where we might want to specify more dependencies:
     # as dependencies for the library being built, and as binary packages
     # to be `cargo install`'d.
@@ -80,19 +111,29 @@ def install_environment(
     if len(lib_deps) > 0:
         _add_dependencies(prefix.path('Cargo.toml'), lib_deps)
 
-    with clean_path_on_failure(directory):
+    language_version = _resolve_version(prefix, version)
+    directory = _envdir(prefix, language_version)
+    with clean_path_on_failure(directory), in_env(prefix, language_version):
+        if language_version != C.DEFAULT:
+            cmd_output_b(
+                'rustup', 'toolchain', 'install',
+                '--no-self-update', '--profile', 'minimal', language_version,
+                cwd=prefix.prefix_dir,
+            )
+
         packages_to_install: Set[Tuple[str, ...]] = {('--path', '.')}
         for cli_dep in cli_deps:
             cli_dep = cli_dep[len('cli:'):]
-            package, _, version = cli_dep.partition(':')
-            if version != '':
-                packages_to_install.add((package, '--version', version))
+            package, _, dep_version = cli_dep.partition(':')
+            if dep_version != '':
+                packages_to_install.add((package, '--version', dep_version))
             else:
                 packages_to_install.add((package,))
 
+        version_flag = _version_flag(language_version)
         for args in packages_to_install:
             cmd_output_b(
-                'cargo', 'install', '--bins', '--root', directory, *args,
+                'cargo', *version_flag, 'install', '--bins', *args,
                 cwd=prefix.prefix_dir,
             )
 
@@ -102,5 +143,6 @@ def run_hook(
         file_args: Sequence[str],
         color: bool,
 ) -> Tuple[int, bytes]:
-    with in_env(hook.prefix):
+    language_version = _resolve_version(hook.prefix, hook.language_version)
+    with in_env(hook.prefix, language_version):
         return helpers.run_xargs(hook, hook.cmd, file_args, color=color)