address comments

Author: R-Palazzo

deepecho/models/_utils.py

@@ -0,0 +1,45 @@
+import sys
+import warnings
+
+import torch
+
+
+def _validate_gpu_parameter(enable_gpu, cuda):
+    if cuda is not None:
+        warnings.warn(
+            '`cuda` parameter is deprecated and will be removed in a future release. '
+            'Please use `enable_gpu` instead.',
+            FutureWarning,
+        )
+        if not enable_gpu:
+            raise ValueError(
+                'Cannot resolve the provided values of `cuda` and `enable_gpu` parameters. '
+                'Please use only `enable_gpu`.'
+            )
+
+        enable_gpu = cuda
+
+    return enable_gpu
+
+
+def _set_device(enable_gpu, device=None):
+    if device:
+        return torch.device(device)
+
+    if enable_gpu:
+        if sys.platform == 'darwin':  # macOS
+            if getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available():
+                device = 'mps'
+            else:
+                device = 'cpu'
+        else:  # Linux/Windows
+            device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    else:
+        device = 'cpu'
+
+    return torch.device(device)
+
+
+def validate_and_set_device(enable_gpu, cuda):
+    enable_gpu = _validate_gpu_parameter(enable_gpu, cuda)
+    return _set_device(enable_gpu)

deepecho/models/basic_gan.py

@@ -1,47 +1,18 @@
 """BasicGAN Model."""
 
 import logging
-import sys
-import warnings
 
 import numpy as np
 import pandas as pd
 import torch
 from tqdm import tqdm
 
+from deepecho.models._utils import validate_and_set_device
 from deepecho.models.base import DeepEcho
 
 LOGGER = logging.getLogger(__name__)
 
 
-def _set_device(enable_gpu, cuda):
-    if cuda is not None:
-        if not enable_gpu:
-            raise ValueError(
-                'Cannot set `cuda` and `enable_gpu` together. Please use only `enable_gpu`.'
-            )
-
-        warnings.warn(
-            '`cuda` parameter is deprecated and will be removed in a future release. '
-            'Please use `enable_gpu` instead.',
-            FutureWarning,
-        )
-        enable_gpu = cuda
-
-    if enable_gpu:
-        if sys.platform == 'darwin':  # macOS
-            if getattr(torch.backends, 'mps', None) and torch.backends.mps.is_available():
-                device = 'mps'
-            else:
-                device = 'cpu'
-        else:  # Linux/Windows
-            device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    else:
-        device = 'cpu'
-
-    return torch.device(device)
-
-
 def _expand_context(data, context):
     return torch.cat(
         [
@@ -170,7 +141,7 @@ class BasicGANModel(DeepEcho):
             Whether to attempt to use GPU for computation.
             Defaults to ``True``.
         cuda (bool):
-            ** Deprecated ** Whether to attempt to use cuda for GPU computation.
+            **Deprecated** Whether to attempt to use cuda for GPU computation.
             If this is False or CUDA is not available, CPU will be used.
         verbose (bool):
             Whether to print progress to console or not.
@@ -201,7 +172,7 @@ def __init__(
         self._dis_lr = dis_lr
         self._latent_size = latent_size
         self._hidden_size = hidden_size
-        self._device = _set_device(enable_gpu, cuda)
+        self._device = validate_and_set_device(enable_gpu, cuda)
         self._enable_gpu = cuda if cuda is not None else enable_gpu
         self._verbose = verbose
 

tests/integration/test_basic_gan.py

@@ -7,44 +7,7 @@
 import pytest
 import torch
 
-from deepecho.models.basic_gan import BasicGANModel, _set_device
-
-
-def test__set_device():
-    """Test the ``_set_device`` method."""
-    # Setup
-    expected_error = re.escape(
-        'Cannot set `cuda` and `enable_gpu` together. Please use only `enable_gpu`.'
-    )
-    expected_warning = re.escape(
-        '`cuda` parameter is deprecated and will be removed in a future release. '
-        'Please use `enable_gpu` instead.'
-    )
-
-    # Run
-    device_1 = _set_device(False, None)
-    device_2 = _set_device(True, None)
-    with pytest.warns(FutureWarning, match=expected_warning):
-        device_3 = _set_device(True, False)
-
-    with pytest.raises(ValueError, match=expected_error):
-        _set_device(False, True)
-
-    # Assert
-    if (
-        sys.platform == 'darwin'
-        and getattr(torch.backends, 'mps', None)
-        and torch.backends.mps.is_available()
-    ):
-        expected_device_2 = torch.device('mps')
-    elif torch.cuda.is_available():
-        expected_device_2 = torch.device('cuda')
-    else:
-        expected_device_2 = torch.device('cpu')
-
-    assert device_1 == torch.device('cpu')
-    assert device_2 == expected_device_2
-    assert device_3 == torch.device('cpu')
+from deepecho.models.basic_gan import BasicGANModel
 
 
 class TestBasicGANModel(unittest.TestCase):

tests/unit/models/test__utils.py

@@ -0,0 +1,80 @@
+"""Test the ``_utils`` module."""
+
+import re
+import sys
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from deepecho.models._utils import _set_device, _validate_gpu_parameter, validate_and_set_device
+
+
+def test__validate_gpu_parameter():
+    """Test the ``_validate_gpu_parameter`` method."""
+    # Setup
+    expected_error = re.escape(
+        'Cannot resolve the provided values of `cuda` and `enable_gpu` parameters. '
+        'Please use only `enable_gpu`.'
+    )
+    expected_warning = re.escape(
+        '`cuda` parameter is deprecated and will be removed in a future release. '
+        'Please use `enable_gpu` instead.'
+    )
+
+    # Run
+    enable_gpu_1 = _validate_gpu_parameter(False, None)
+    enable_gpu_2 = _validate_gpu_parameter(True, None)
+    with pytest.warns(FutureWarning, match=expected_warning):
+        enable_gpu_3 = _validate_gpu_parameter(True, False)
+
+    with pytest.raises(ValueError, match=expected_error):
+        _validate_gpu_parameter(False, True)
+
+    # Assert
+    assert enable_gpu_1 is False
+    assert enable_gpu_2 is True
+    assert enable_gpu_3 is False
+
+
+def test__set_device():
+    """Test the ``_set_device`` method."""
+    # Run
+    device_1 = _set_device(False)
+    device_2 = _set_device(True)
+    device_3 = _set_device(True, 'cpu')
+    device_4 = _set_device(enable_gpu=False, device='cpu')
+
+    # Assert
+    if (
+        sys.platform == 'darwin'
+        and getattr(torch.backends, 'mps', None)
+        and torch.backends.mps.is_available()
+    ):
+        expected_device_2 = torch.device('mps')
+    elif torch.cuda.is_available():
+        expected_device_2 = torch.device('cuda')
+    else:
+        expected_device_2 = torch.device('cpu')
+
+    assert device_1 == torch.device('cpu')
+    assert device_2 == expected_device_2
+    assert device_3 == torch.device('cpu')
+    assert device_4 == torch.device('cpu')
+
+
+@patch('deepecho.models._utils._set_device')
+@patch('deepecho.models._utils._validate_gpu_parameter')
+def test_validate_and_set_device(mock_validate, mock_set_device):
+    """Test the ``validate_and_set_device`` method."""
+    # Setup
+    mock_validate.return_value = True
+    mock_set_device.return_value = torch.device('cuda')
+
+    # Run
+    device = validate_and_set_device(enable_gpu=True, cuda=None)
+
+    # Assert
+    mock_validate.assert_called_once_with(True, None)
+    mock_set_device.assert_called_once_with(True)
+    assert device == torch.device('cuda')