From ce4294215e5d2b3aca06d761d466d44e528cc4c4 Mon Sep 17 00:00:00 2001
From: Erwan Rouchet <rouchet@teklia.com>
Date: Fri, 29 Nov 2019 09:54:44 +0100
Subject: [PATCH] Add DeviceRequest type

---
 docker/types/__init__.py   |  3 +-
 docker/types/containers.py | 98 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 100 insertions(+), 1 deletion(-)

diff --git a/docker/types/__init__.py b/docker/types/__init__.py
index 5db330e284..91401bed41 100644
--- a/docker/types/__init__.py
+++ b/docker/types/__init__.py
@@ -1,5 +1,6 @@
 # flake8: noqa
-from .containers import ContainerConfig, HostConfig, LogConfig, Ulimit
+from .containers import \
+    ContainerConfig, HostConfig, LogConfig, Ulimit, DeviceRequest
 from .daemon import CancellableStream
 from .healthcheck import Healthcheck
 from .networks import EndpointConfig, IPAMConfig, IPAMPool, NetworkingConfig
diff --git a/docker/types/containers.py b/docker/types/containers.py
index fd8cab4979..a1a5b177a1 100644
--- a/docker/types/containers.py
+++ b/docker/types/containers.py
@@ -154,6 +154,104 @@ def hard(self, value):
         self['Hard'] = value
 
 
+class DeviceRequest(DictType):
+    """
+    Create a device request to be used with
+    :py:meth:`~docker.api.container.ContainerApiMixin.create_host_config`.
+
+    Args:
+
+        driver (str): Which driver to use for this device. Optional.
+        count (int): Number or devices to request. Optional.
+            Set to -1 to request all available devices.
+        device_ids (list): List of strings for device IDs. Optional.
+            Set either ``count`` or ``device_ids``.
+        capabilities (list): List of lists of strings to request
+            capabilities. Optional. The global list acts like an OR,
+            and the sub-lists are AND. The driver will try to satisfy
+            one of the sub-lists.
+            Available capabilities for the ``nvidia`` driver can be found
+            `here <https://github.com/NVIDIA/nvidia-container-runtime>`_.
+        options (dict): Driver-specific options. Optional.
+    """
+
+    def __init__(self, **kwargs):
+        driver = kwargs.get('driver', kwargs.get('Driver'))
+        count = kwargs.get('count', kwargs.get('Count'))
+        device_ids = kwargs.get('device_ids', kwargs.get('DeviceIDs'))
+        capabilities = kwargs.get('capabilities', kwargs.get('Capabilities'))
+        options = kwargs.get('options', kwargs.get('Options'))
+
+        if driver is None:
+            driver = ''
+        elif not isinstance(driver, six.string_types):
+            raise ValueError('DeviceRequest.driver must be a string')
+        if count is None:
+            count = 0
+        elif not isinstance(count, int):
+            raise ValueError('DeviceRequest.count must be an integer')
+        if device_ids is None:
+            device_ids = []
+        elif not isinstance(device_ids, list):
+            raise ValueError('DeviceRequest.device_ids must be a list')
+        if capabilities is None:
+            capabilities = []
+        elif not isinstance(capabilities, list):
+            raise ValueError('DeviceRequest.capabilities must be a list')
+        if options is None:
+            options = {}
+        elif not isinstance(options, dict):
+            raise ValueError('DeviceRequest.options must be a dict')
+
+        super(DeviceRequest, self).__init__({
+            'Driver': driver,
+            'Count': count,
+            'DeviceIDs': device_ids,
+            'Capabilities': capabilities,
+            'Options': options
+        })
+
+    @property
+    def driver(self):
+        return self['Driver']
+
+    @driver.setter
+    def driver(self, value):
+        self['Driver'] = value
+
+    @property
+    def count(self):
+        return self['Count']
+
+    @count.setter
+    def count(self, value):
+        self['Count'] = value
+
+    @property
+    def device_ids(self):
+        return self['DeviceIDs']
+
+    @device_ids.setter
+    def device_ids(self, value):
+        self['DeviceIDs'] = value
+
+    @property
+    def capabilities(self):
+        return self['Capabilities']
+
+    @capabilities.setter
+    def capabilities(self, value):
+        self['Capabilities'] = value
+
+    @property
+    def options(self):
+        return self['Options']
+
+    @options.setter
+    def options(self, value):
+        self['Options'] = value
+
+
 class HostConfig(dict):
     def __init__(self, version, binds=None, port_bindings=None,
                  lxc_conf=None, publish_all_ports=False, links=None,