use staticmethod for new types

Author: isabelizimm

.github/workflows/tests.yml

@@ -64,3 +64,21 @@ jobs:
       - name: Run tests
         run: |
           pytest vetiver -m 'rsc_test'
+
+  test-no-torch:
+    name: "Test no-torch"
+    runs-on: ubuntu-latest
+    if: ${{ !github.event.pull_request.head.repo.fork }}
+    steps:
+      - uses: actions/checkout@v2
+      - uses: actions/setup-python@v2
+        with:
+          python-version: 3.8
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e .[dev]
+
+      - name: Run tests
+        run: |
+          pytest vetiver/tests/test_sklearn.py

vetiver/handlers/base.py

@@ -68,11 +68,11 @@ class VetiverHandler:
     def __init_subclass__(cls, **kwargs):
         super().__init_subclass__(**kwargs)
         with suppress(AttributeError, NameError):
-            create_handler.register(cls.base_class, cls)
+            create_handler.register(cls.base_class(), cls)
 
-    def __new__(cls, value=None):
-        implementation_cls = create_handler.registry[type(value)]
-        return super().__new__(implementation_cls)
+    # def __new__(cls, value=None):
+    #     implementation_cls = create_handler.registry[type(value)]
+    #     return super().__new__(implementation_cls)
 
     def __init__(self, model, ptype_data):
         self.model = model

vetiver/handlers/sklearn.py

@@ -14,7 +14,7 @@ class SKLearnHandler(VetiverHandler):
         a trained sklearn model
     """
 
-    base_class = sklearn.base.BaseEstimator
+    base_class = staticmethod(lambda: sklearn.base.BaseEstimator)
 
     def __init__(self, model, ptype_data):
         super().__init__(model, ptype_data)
@@ -54,7 +54,7 @@ def handler_predict(self, input_data, check_ptype):
             Prediction from model
         """
 
-        if check_ptype == True:
+        if check_ptype:
             if isinstance(input_data, pd.DataFrame):
                 prediction = self.model.predict(input_data)
             else:

vetiver/handlers/torch.py

@@ -19,7 +19,7 @@ class TorchHandler(VetiverHandler):
         a trained torch model
     """
 
-    base_class = torch.nn.Module
+    base_class = staticmethod(lambda: torch.nn.Module)
 
     def __init__(self, model, ptype_data):
         super().__init__(model, ptype_data)
@@ -59,7 +59,7 @@ def handler_predict(self, input_data, check_ptype):
             Prediction from model
         """
         if torch_exists:
-            if check_ptype == True:
+            if check_ptype:
                 input_data = np.array(input_data, dtype=np.array(self.ptype_data).dtype)
                 prediction = self.model(torch.from_numpy(input_data))
 

vetiver/tests/test_pytorch.py

@@ -1,19 +1,19 @@
 from vetiver.vetiver_model import VetiverModel
 from vetiver import VetiverAPI
 from fastapi.testclient import TestClient
-
-import torch
-import torch.nn as nn
+import pytest
 import numpy as np
 
+torch = pytest.importorskip("torch")
+
+import torch  # noqa
+
 
 def _build_torch_v():
 
-    # Hyper-parameters
     input_size = 1
     output_size = 1
 
-    # Toy dataset
     x_train = np.array(
         [
             [3.3],
@@ -34,8 +34,8 @@ def _build_torch_v():
         ],
         dtype=np.float32,
     )
-    # Linear regression model
-    torch_model = nn.Linear(input_size, output_size)
+
+    torch_model = torch.nn.Linear(input_size, output_size)
     return x_train, torch_model