Add type annotation for get_driver and set_driver.

Kami · Kami · commit f5a201637f51 · 2019-07-10T11:07:28.000+02:00
diff --git a/libcloud/compute/providers.py b/libcloud/compute/providers.py
@@ -19,6 +19,7 @@
 from __future__ import absolute_import
 
 from typing import Type
+from typing import Union
 from types import ModuleType
 
 from libcloud.compute.types import Provider
@@ -163,14 +164,14 @@
 
 
 def get_driver(provider):
-    # type: (Provider) -> Type[NodeDriver]
+    # type: (Union[Provider, str]) -> Type[NodeDriver]
     deprecated_constants = OLD_CONSTANT_TO_NEW_MAPPING
     return _get_provider_driver(drivers=DRIVERS, provider=provider,
                                 deprecated_providers=DEPRECATED_DRIVERS,
                                 deprecated_constants=deprecated_constants)
 
 
 def set_driver(provider, module, klass):
-    # type: (Provider, ModuleType, type) -> Type[NodeDriver]
+    # type: (Union[Provider, str], ModuleType, type) -> Type[NodeDriver]
     return _set_provider_driver(drivers=DRIVERS, provider=provider,
                                 module=module, klass=klass)
