From 91111ab7d8e6b4c3d510b7c044c58f915f5db177 Mon Sep 17 00:00:00 2001
From: nat <n@natn.me>
Date: Thu, 28 May 2020 19:28:53 +0200
Subject: [PATCH] Make plugin import errors more helpful

This addresses an issue where if the user happens to have black installed in
their environment but not the other dependencies when running the protoc
plugin then the resulting import error (No module named 'google') is not very
helpful.
---
 betterproto/plugin.py              | 36 +++++++++++++++++-------------
 betterproto/tests/inputs/config.py |  2 +-
 2 files changed, 21 insertions(+), 17 deletions(-)

diff --git a/betterproto/plugin.py b/betterproto/plugin.py
index 8eecef2..c843a6a 100755
--- a/betterproto/plugin.py
+++ b/betterproto/plugin.py
@@ -1,33 +1,37 @@
 #!/usr/bin/env python
 
+from collections import defaultdict
 import itertools
 import os.path
+import stringcase
 import sys
 import textwrap
-from collections import defaultdict
 from typing import Dict, List, Optional, Type
+from betterproto.casing import safe_snake_case
 
 try:
+    # betterproto[compiler] specific dependencies
     import black
-except ImportError:
+    from google.protobuf.compiler import plugin_pb2 as plugin
+    from google.protobuf.descriptor_pb2 import (
+        DescriptorProto,
+        EnumDescriptorProto,
+        FieldDescriptorProto,
+    )
+    import google.protobuf.wrappers_pb2 as google_wrappers
+    import jinja2
+except ImportError as err:
+    missing_import = err.args[0][17:-1]
     print(
-        "Unable to import `black` formatter. Did you install the compiler feature with `pip install betterproto[compiler]`?"
+        "\033[31m"
+        f"Unable to import `{missing_import}` from betterproto plugin! "
+        "Please ensure that you've installed betterproto as "
+        '`pip install "betterproto[compiler]"` so that compiler dependencies '
+        "are included."
+        "\033[0m"
     )
     raise SystemExit(1)
 
-import jinja2
-import stringcase
-
-from google.protobuf.compiler import plugin_pb2 as plugin
-from google.protobuf.descriptor_pb2 import (
-    DescriptorProto,
-    EnumDescriptorProto,
-    FieldDescriptorProto,
-)
-
-from betterproto.casing import safe_snake_case
-
-import google.protobuf.wrappers_pb2 as google_wrappers
 
 WRAPPER_TYPES: Dict[str, Optional[Type]] = defaultdict(
     lambda: None,
diff --git a/betterproto/tests/inputs/config.py b/betterproto/tests/inputs/config.py
index cf4c996..1f276ba 100644
--- a/betterproto/tests/inputs/config.py
+++ b/betterproto/tests/inputs/config.py
@@ -11,7 +11,7 @@ tests = {
     "googletypes_service_returns_empty",  # 9
     "casing_message_field_uppercase",  # 11
     "namespace_keywords",  # 70
-    "namespace_builtin_types"  # 53
+    "namespace_builtin_types",  # 53
 }
 
 services = {