Support other lookups on ModelInstance

This makes `ModelInstance` accept a keyword argument `lookup` (that
defaults to 'pk') which allows to specify which lookup to use on the
given `QuerySet`.

Author: ulope

djclick/params.py

@@ -4,7 +4,7 @@
 
 
 class ModelInstance(click.ParamType):
-    def __init__(self, qs):
+    def __init__(self, qs, lookup='pk'):
         from django.db import models
 
         if isinstance(qs, type) and issubclass(qs, models.Model):
@@ -14,13 +14,17 @@ def __init__(self, qs):
             qs.model._meta.app_label,
             qs.model.__name__,
         )
+        self.lookup = lookup
 
     def convert(self, value, param, ctx):
+        if value is None:
+            return super(ModelInstance, self).convert(value, param, ctx)
         try:
-            return self.qs.get(pk=value)
+            return self.qs.get(**{self.lookup: value})
         except ObjectDoesNotExist:
             pass
         # call `fail` outside of exception context to avoid nested exception
         # handling on Python 3
-        msg = 'could not find {} with pk={}'.format(self.name, value)
+        msg = 'could not find {s.name} with {s.lookup}={value}'.format(
+            s=self, value=value)
         self.fail(msg, param, ctx)

djclick/test/test_params.py

@@ -1,5 +1,5 @@
-import pytest
 from click.exceptions import BadParameter
+import pytest
 
 from djclick import params
 
@@ -18,17 +18,37 @@ def test_modelinstance_init():
 
 
 @pytest.mark.django_db
-def test_convert_ok(call_command):
+@pytest.mark.parametrize(
+    ('arg', 'value'),
+    (
+        ('--pk', '99'),
+        ('--slug', 'test'),
+        ('--endswith', 'st'),
+    )
+)
+def test_convert_ok(call_command, arg, value):
     from testapp.models import DummyModel
 
-    DummyModel.objects.create()
-    assert call_command('modelcmd', '1').stdout == b'1'
+    DummyModel.objects.create(slug='test')
+    expected = b'<DummyModel: 1>'
+
+    assert call_command('modelcmd', '--pk', '1').stdout == expected
+    assert call_command('modelcmd', '--slug', 'test').stdout == expected
+    assert call_command('modelcmd', '--endswith', 'test').stdout == expected
 
 
 @pytest.mark.django_db
-def test_convert_fail(call_command):
+@pytest.mark.parametrize(
+    ('args', 'error_message'),
+    (
+        (('--pk', '99'), "pk=99"),
+        (('--slug', 'doesnotexist'), "slug=doesnotexist"),
+    )
+)
+def test_convert_fail(call_command, args, error_message):
     with pytest.raises(BadParameter) as e:
-        call_command('modelcmd', '999')
+        call_command('modelcmd', *args)
     # Use `.endswith()` because of differences between CPython and pypy
-    assert str(e).endswith('BadParameter: could not find '
-                           'testapp.DummyModel with pk=999')
+    assert str(e).endswith(
+        'BadParameter: could not find testapp.DummyModel with {}'.format(
+            error_message))

djclick/test/testprj/testapp/management/commands/modelcmd.py

@@ -5,7 +5,14 @@
 
 
 @click.command()
-@click.argument('instance', type=ModelInstance(DummyModel.objects.all()))
-def command(instance):
-    # Just print some things which shall not be found in the output
-    click.echo(instance, nl=False)
+@click.option('--pk', type=ModelInstance(DummyModel))
+@click.option('--slug', type=ModelInstance(DummyModel, lookup="slug"))
+@click.option('--endswith', type=ModelInstance(DummyModel,
+                                               lookup="slug__endswith"))
+def command(pk, slug, endswith):
+    if pk:
+        click.echo(repr(pk), nl=False)
+    if slug:
+        click.echo(repr(slug), nl=False)
+    if endswith:
+        click.echo(repr(endswith), nl=False)

djclick/test/testprj/testapp/models.py

@@ -6,5 +6,7 @@
 
 @python_2_unicode_compatible
 class DummyModel(models.Model):
+    slug = models.CharField(max_length=50)
+
     def __str__(self):
         return str(self.id)