From 7f6b4be05e55876184fa92eca2886ce1765bb655 Mon Sep 17 00:00:00 2001
From: mike0sv <mike0sv@gmail.com>
Date: Wed, 18 May 2022 08:12:57 +0300
Subject: [PATCH] better errors closes #230

---
 mlem/api/commands.py         |  4 +---
 mlem/contrib/docker/base.py  |  2 ++
 mlem/contrib/docker/utils.py | 14 ++++++++++++++
 mlem/core/base.py            |  8 ++++++++
 mlem/core/errors.py          |  7 +++++++
 5 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/mlem/api/commands.py b/mlem/api/commands.py
index e01168bc..6a1c718d 100644
--- a/mlem/api/commands.py
+++ b/mlem/api/commands.py
@@ -394,9 +394,7 @@ def import_object(
     echo(EMOJI_LOAD + f"Importing object from {loc.uri_repr}")
     if type_ is not None:
         type_, modifier = parse_import_type_modifier(type_)
-        if type_ not in ImportHook.__type_map__:
-            raise ValueError(f"Unknown import type {type_}")
-        meta = ImportHook.__type_map__[type_].process(
+        meta = ImportHook.load_type(type_).process(
             loc, copy_data=copy_data, modifier=modifier
         )
     else:
diff --git a/mlem/contrib/docker/base.py b/mlem/contrib/docker/base.py
index 0087ebb1..87936b07 100644
--- a/mlem/contrib/docker/base.py
+++ b/mlem/contrib/docker/base.py
@@ -16,6 +16,7 @@
     create_docker_client,
     image_exists_at_dockerhub,
     print_docker_logs,
+    wrap_docker_error,
 )
 from mlem.core.base import MlemABC
 from mlem.core.errors import DeploymentError
@@ -336,6 +337,7 @@ def package(self, obj: MlemModel) -> DockerImage:
 
             return self.build(tempdir)
 
+    @wrap_docker_error
     def build(self, context_dir: str) -> DockerImage:
         tag = self.image.uri
         logger.debug("Building docker image %s from %s...", tag, context_dir)
diff --git a/mlem/contrib/docker/utils.py b/mlem/contrib/docker/utils.py
index d9204428..77c1543f 100644
--- a/mlem/contrib/docker/utils.py
+++ b/mlem/contrib/docker/utils.py
@@ -4,6 +4,7 @@
 import re
 import time
 from contextlib import contextmanager
+from functools import wraps
 from threading import Lock
 from typing import Any, Iterator, Tuple, Union
 
@@ -13,6 +14,8 @@
 from docker.models.images import Image
 from docker.utils.json_stream import json_stream
 
+from mlem.core.errors import MlemError
+
 logger = logging.getLogger(__name__)
 
 
@@ -124,6 +127,17 @@ def repository_tags_at_dockerhub(repo):
     return {tag["name"] for tag in resp.json()}
 
 
+def wrap_docker_error(f):
+    @wraps(f)
+    def inner(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except DockerException as e:
+            raise MlemError(f"Error calling docker: {e}") from e
+
+    return inner
+
+
 # Copyright 2019 Zyfra
 # Copyright 2021 Iterative
 #
diff --git a/mlem/core/base.py b/mlem/core/base.py
index ae6e397b..2e337341 100644
--- a/mlem/core/base.py
+++ b/mlem/core/base.py
@@ -6,6 +6,7 @@
 from typing_extensions import Literal
 from yaml import safe_load
 
+from mlem.core.errors import UnknownImplementation
 from mlem.polydantic import PolyModel
 from mlem.utils.importing import import_string
 from mlem.utils.path import make_posix
@@ -102,6 +103,13 @@ def non_abstract_subtypes(cls: Type[MT]) -> Dict[str, Type["MT"]]:
             and v is not cls
         }
 
+    @classmethod
+    def load_type(cls, type_name: str):
+        try:
+            return cls.__resolve_subtype__(type_name)
+        except ValueError as e:
+            raise UnknownImplementation(type_name, cls.abs_name) from e
+
 
 def set_or_replace(obj: dict, key: str, value: Any, subkey: str = "type"):
     if key in obj:
diff --git a/mlem/core/errors.py b/mlem/core/errors.py
index 5b854c88..b4aa65a8 100644
--- a/mlem/core/errors.py
+++ b/mlem/core/errors.py
@@ -130,3 +130,10 @@ def __init__(self, wrong, missing, fix):
         super().__init__(
             f"Wrong requirements: {self.wrong} {self.missing}\nTo fix it, run `{fix}`"
         )
+
+
+class UnknownImplementation(MlemError):
+    def __init__(self, type_name: str, abs_name: str):
+        self.abs_name = abs_name
+        self.type_name = type_name
+        super().__init__(f"Unknown {abs_name} implementation: {type_name}")