From 62dddd5821312aca4a6001c7b84be0bee137fc48 Mon Sep 17 00:00:00 2001
From: Divy Jain <dkj@somaiya.edu>
Date: Wed, 20 Jan 2021 03:07:05 +0530
Subject: [PATCH] #1149: Add python_module_error rule (#1151)

---
 README.md                               |  1 +
 tests/rules/test_python_module_error.py | 63 +++++++++++++++++++++++++
 thefuck/rules/python_module_error.py    | 13 +++++
 3 files changed, 77 insertions(+)
 create mode 100644 tests/rules/test_python_module_error.py
 create mode 100644 thefuck/rules/python_module_error.py

diff --git a/README.md b/README.md
index 5ac8b83e..8aa2294d 100644
--- a/README.md
+++ b/README.md
@@ -281,6 +281,7 @@ following rules are enabled by default:
 * `pyenv_no_such_command` &ndash; fixes wrong pyenv commands like `pyenv isntall` or `pyenv list`;
 * `python_command` &ndash; prepends `python` when you try to run non-executable/without `./` python script;
 * `python_execute` &ndash; appends missing `.py` when executing Python files;
+* `python_module_error` &ndash; fixes ModuleNotFoundError by trying to `pip install` that module;
 * `quotation_marks` &ndash; fixes uneven usage of `'` and `"` when containing args';
 * `path_from_history` &ndash; replaces not found path with similar absolute path from history;
 * `react_native_command_unrecognized` &ndash; fixes unrecognized `react-native` commands;
diff --git a/tests/rules/test_python_module_error.py b/tests/rules/test_python_module_error.py
new file mode 100644
index 00000000..838f9561
--- /dev/null
+++ b/tests/rules/test_python_module_error.py
@@ -0,0 +1,63 @@
+import pytest
+
+from thefuck.rules.python_module_error import get_new_command, match
+from thefuck.types import Command
+
+
+@pytest.fixture
+def module_error_output(filename, module_name):
+    return """Traceback (most recent call last):
+  File "{0}", line 1, in <module>
+    import {1}
+ModuleNotFoundError: No module named '{1}'""".format(
+        filename, module_name
+    )
+
+
+@pytest.mark.parametrize(
+    "test",
+    [
+        Command("python hello_world.py", "Hello World"),
+        Command(
+            "./hello_world.py",
+            """Traceback (most recent call last):
+  File "hello_world.py", line 1, in <module>
+    pritn("Hello World")
+NameError: name 'pritn' is not defined""",
+        ),
+    ],
+)
+def test_not_match(test):
+    assert not match(test)
+
+
+positive_tests = [
+    (
+        "python some_script.py",
+        "some_script.py",
+        "more_itertools",
+        "pip install more_itertools && python some_script.py",
+    ),
+    (
+        "./some_other_script.py",
+        "some_other_script.py",
+        "a_module",
+        "pip install a_module && ./some_other_script.py",
+    ),
+]
+
+
+@pytest.mark.parametrize(
+    "script, filename, module_name, corrected_script", positive_tests
+)
+def test_match(script, filename, module_name, corrected_script, module_error_output):
+    assert match(Command(script, module_error_output))
+
+
+@pytest.mark.parametrize(
+    "script, filename, module_name, corrected_script", positive_tests
+)
+def test_get_new_command(
+    script, filename, module_name, corrected_script, module_error_output
+):
+    assert get_new_command(Command(script, module_error_output)) == corrected_script
diff --git a/thefuck/rules/python_module_error.py b/thefuck/rules/python_module_error.py
new file mode 100644
index 00000000..4696d63b
--- /dev/null
+++ b/thefuck/rules/python_module_error.py
@@ -0,0 +1,13 @@
+import re
+from thefuck.shells import shell
+
+MISSING_MODULE = r"ModuleNotFoundError: No module named '([^']+)'"
+
+
+def match(command):
+    return "ModuleNotFoundError: No module named '" in command.output
+
+
+def get_new_command(command):
+    missing_module = re.findall(MISSING_MODULE, command.output)[0]
+    return shell.and_("pip install {}".format(missing_module), command.script)