Repository: HypothesisWorks/hypothesis

Issue:
register custom strategies for NewTypes
Hey guys,

First off thanks for the great work with hypothesis! It's an extremely useful testing tool.

I read from previous open issues that you have added support for `NewType` in `from_type`. It is however currently not possible to register a custom strategy for a `NewType`. I get the same error as you fixed in `from_type`:
`hypothesis.errors.InvalidArgument: custom_type=%r must be a type`

I use hypothesis v. 4.14.2.




Discussion:
What version of Python?  Can you provide example code?
I use python 3.7. 

An Example:

```
Name = NewType('Name', str)

register_type_strategy(Name, characters(whitelist_categories=('L',)))
```
I just made a pull request for this. I hope that's ok:
https://github.com/HypothesisWorks/hypothesis/pull/1914

I duplicated the code that fetched the strategy from the _global_type_lookup. Not very DRY but I was unsure of where to put it.

I also had a test, which I did not know where to add. It very naive - I would have rather tested against the type of the strategy returned but from_type returns a deferred strategy which I didn't know how to forward:

```
def test_new_type():

    # type not registered
    MyNewType = typing.NewType('MyNewType', str)
    example = st.from_type(MyNewType).example()
    assert isinstance(example, str)

    # register type and check again
    strategy = st.characters(whitelist_categories='L')
    st.register_type_strategy(MyNewType, strategy)
    example = st.from_type(MyNewType).example()
    assert isinstance(example, str) and len(example) <= 1
```



Resolving patch:
diff --git a/hypothesis-python/src/hypothesis/_strategies.py b/hypothesis-python/src/hypothesis/_strategies.py
--- a/hypothesis-python/src/hypothesis/_strategies.py
+++ b/hypothesis-python/src/hypothesis/_strategies.py
@@ -26,7 +26,7 @@
 from decimal import Context, Decimal, localcontext
 from fractions import Fraction
 from functools import reduce
-from inspect import isclass, isfunction
+from inspect import isclass
 from uuid import UUID
 
 import attr
@@ -67,6 +67,7 @@
 from hypothesis.internal.reflection import (
     define_function_signature,
     is_typed_named_tuple,
+    nicerepr,
     proxies,
     required_args,
 )
@@ -1287,17 +1288,30 @@ def from_type(thing):
     # refactoring it's hard to do without creating circular imports.
     from hypothesis.searchstrategy import types
 
+    def as_strategy(strat_or_callable, thing):
+        # User-provided strategies need some validation, and callables even more
+        # of it.  We do this in three places, hence the helper function
+        if not isinstance(strat_or_callable, SearchStrategy):
+            assert callable(strat_or_callable)  # Validated in register_type_strategy
+            strategy = strat_or_callable(thing)
+        else:
+            strategy = strat_or_callable
+        if not isinstance(strategy, SearchStrategy):
+            raise ResolutionFailed(
+                "Error: %s was registered for %r, but returned non-strategy %r"
+                % (thing, nicerepr(strat_or_callable), strategy)
+            )
+        if strategy.is_empty:
+            raise ResolutionFailed("Error: %r resolved to an empty strategy" % (thing,))
+        return strategy
+
     if typing is not None:  # pragma: no branch
         if not isinstance(thing, type):
-            # At runtime, `typing.NewType` returns an identity function rather
-            # than an actual type, but we can check that for a possible match
-            # and then read the magic attribute to unwrap it.
-            if (
-                hasattr(thing, "__supertype__")
-                and hasattr(typing, "NewType")
-                and isfunction(thing)
-                and getattr(thing, "__module__", 0) == "typing"
-            ):
+            if types.is_a_new_type(thing):
+                # Check if we have an explicitly registered strategy for this thing,
+                # resolve it so, and otherwise resolve as for the base type.
+                if thing in types._global_type_lookup:
+                    return as_strategy(types._global_type_lookup[thing], thing)
                 return from_type(thing.__supertype__)
             # Under Python 3.6, Unions are not instances of `type` - but we
             # still want to resolve them!
@@ -1321,12 +1335,7 @@ def from_type(thing):
     # lookup may be a strategy or a function from type -> strategy; and we
     # convert empty results into an explicit error.
     if thing in types._global_type_lookup:
-        strategy = types._global_type_lookup[thing]
-        if not isinstance(strategy, SearchStrategy):
-            strategy = strategy(thing)  # type: ignore
-        if strategy.is_empty:  # type: ignore
-            raise ResolutionFailed("Error: %r resolved to an empty strategy" % (thing,))
-        return strategy  # type: ignore
+        return as_strategy(types._global_type_lookup[thing], thing)
     # If there's no explicitly registered strategy, maybe a subtype of thing
     # is registered - if so, we can resolve it to the subclass strategy.
     # We'll start by checking if thing is from from the typing module,
@@ -1340,19 +1349,13 @@ def from_type(thing):
     # type.  For example, `Number -> integers() | floats()`, but bools() is
     # not included because bool is a subclass of int as well as Number.
     strategies = [
-        v if isinstance(v, SearchStrategy) else v(thing)  # type: ignore
+        as_strategy(v, thing)
         for k, v in types._global_type_lookup.items()
         if isinstance(k, type)
         and issubclass(k, thing)
         and sum(types.try_issubclass(k, typ) for typ in types._global_type_lookup) == 1
     ]
-    empty = ", ".join(repr(s) for s in strategies if s.is_empty)
-    if empty:
-        raise ResolutionFailed(
-            "Could not resolve %s to a strategy; consider using "
-            "register_type_strategy" % empty
-        )
-    elif strategies:
+    if strategies:
         return one_of(strategies)
     # If we don't have a strategy registered for this type or any subtype, we
     # may be able to fall back on type annotations.
diff --git a/hypothesis-python/src/hypothesis/searchstrategy/types.py b/hypothesis-python/src/hypothesis/searchstrategy/types.py
--- a/hypothesis-python/src/hypothesis/searchstrategy/types.py
+++ b/hypothesis-python/src/hypothesis/searchstrategy/types.py
@@ -22,6 +22,7 @@
 import decimal
 import fractions
 import functools
+import inspect
 import io
 import numbers
 import uuid
@@ -61,9 +62,24 @@ def try_issubclass(thing, superclass):
         return False
 
 
+def is_a_new_type(thing):
+    # At runtime, `typing.NewType` returns an identity function rather
+    # than an actual type, but we can check whether that thing matches.
+    return (
+        hasattr(thing, "__supertype__")
+        and hasattr(typing, "NewType")
+        and inspect.isfunction(thing)
+        and getattr(thing, "__module__", None) == "typing"
+    )
+
+
 def is_a_type(thing):
     """Return True if thing is a type or a generic type like thing."""
-    return isinstance(thing, type) or isinstance(thing, typing_root_type)
+    return (
+        isinstance(thing, type)
+        or isinstance(thing, typing_root_type)
+        or is_a_new_type(thing)
+    )
 
 
 def from_typing_type(thing):


Test patch:
diff --git a/hypothesis-python/tests/cover/test_type_lookup.py b/hypothesis-python/tests/cover/test_type_lookup.py
--- a/hypothesis-python/tests/cover/test_type_lookup.py
+++ b/hypothesis-python/tests/cover/test_type_lookup.py
@@ -116,10 +116,13 @@ def __init__(self, arg):
         pass
 
 
-def test_custom_type_resolution():
+def test_custom_type_resolution_fails_without_registering():
     fails = st.from_type(UnknownType)
     with pytest.raises(ResolutionFailed):
         fails.example()
+
+
+def test_custom_type_resolution():
     sentinel = object()
     try:
         st.register_type_strategy(UnknownType, st.just(sentinel))
@@ -129,9 +132,29 @@ def test_custom_type_resolution():
     finally:
         types._global_type_lookup.pop(UnknownType)
         st.from_type.__clear_cache()
-    fails = st.from_type(UnknownType)
-    with pytest.raises(ResolutionFailed):
-        fails.example()
+        assert UnknownType not in types._global_type_lookup
+
+
+def test_custom_type_resolution_with_function():
+    sentinel = object()
+    try:
+        st.register_type_strategy(UnknownType, lambda _: st.just(sentinel))
+        assert st.from_type(UnknownType).example() is sentinel
+        assert st.from_type(ParentUnknownType).example() is sentinel
+    finally:
+        types._global_type_lookup.pop(UnknownType)
+        st.from_type.__clear_cache()
+
+
+def test_custom_type_resolution_with_function_non_strategy():
+    try:
+        st.register_type_strategy(UnknownType, lambda _: None)
+        with pytest.raises(ResolutionFailed):
+            st.from_type(UnknownType).example()
+        with pytest.raises(ResolutionFailed):
+            st.from_type(ParentUnknownType).example()
+    finally:
+        types._global_type_lookup.pop(UnknownType)
 
 
 def test_errors_if_generic_resolves_empty():
diff --git a/hypothesis-python/tests/py3/test_lookup.py b/hypothesis-python/tests/py3/test_lookup.py
--- a/hypothesis-python/tests/py3/test_lookup.py
+++ b/hypothesis-python/tests/py3/test_lookup.py
@@ -467,3 +467,9 @@ def test_resolving_recursive_type():
 def test_resolves_empty_Tuple_issue_1583_regression(ex):
     # See e.g. https://github.com/python/mypy/commit/71332d58
     assert ex == ()
+
+
+def test_can_register_NewType():
+    Name = typing.NewType("Name", str)
+    st.register_type_strategy(Name, st.just("Eric Idle"))
+    assert st.from_type(Name).example() == "Eric Idle"