Coverage for haystack/lazy_imports.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from types import TracebackType 

6 

7from lazy_imports.try_import import _DeferredImportExceptionContextManager 

8 

9DEFAULT_IMPORT_ERROR_MSG = "Try 'pip install {}'" 

10 

11 

12class LazyImport(_DeferredImportExceptionContextManager): 

13 """ 

14 A context manager that provides controlled handling of import errors. 

15 

16 It adds the possibility to customize the error messages. 

17 

18 NOTE: Despite its name, this class does not delay the actual import operation. 

19 For installed modules: executes the import immediately. 

20 For uninstalled modules: captures the error and defers it until check() is called. 

21 """ 

22 

23 def __init__(self, message: str = DEFAULT_IMPORT_ERROR_MSG) -> None: 

24 super().__init__() 

25 self.import_error_msg = message 

26 

27 def __exit__( 

28 self, exc_type: type[Exception] | None, exc_value: Exception | None, traceback: TracebackType | None 

29 ) -> bool | None: 

30 """ 

31 Exit the context manager. 

32 

33 Args: 

34 exc_type: 

35 Raised exception type. :obj:`None` if nothing is raised. 

36 exc_value: 

37 Raised exception object. :obj:`None` if nothing is raised. 

38 traceback: 

39 Associated traceback. :obj:`None` if nothing is raised. 

40 

41 Returns: 

42 :obj:`None` if nothing is deferred, otherwise :obj:`True`. 

43 :obj:`True` will suppress any exceptions avoiding them from propagating. 

44 

45 """ 

46 if isinstance(exc_value, ImportError): 

47 message = ( 

48 f"Haystack failed to import the optional dependency '{exc_value.name}'. " 

49 f"{self.import_error_msg.format(exc_value.name)}. Original error: {exc_value}" 

50 ) 

51 self._deferred = (exc_value, message) 

52 return True 

53 return None