Add assert_not_both_not_none (fixes #393)

chex/_src/asserts.py

@@ -248,6 +248,20 @@ def assert_not_both_none(first: Any, second: Any) -> None:
     raise AssertionError(
         "At least one of the arguments must be different from `None`.")
 
+@_static_assertion
+def assert_not_both_not_none(first: Any, second: Any) -> None:
+  """Checks that not both arguments are non-None.
+
+  Args:
+    first: A first object.
+    second: A second object.
+
+  Raises:
+    AssertionError: If both ``first`` and ``second`` are not None.
+  """
+  if first is not None and second is not None:
+    raise AssertionError(
+        "At most one of the arguments may be different from `None`.")
 
 @_static_assertion
 def assert_exactly_one_is_none(first: Any, second: Any) -> None:

chex/_src/asserts_test.py

@@ -1860,7 +1860,15 @@ def test_assert_equal_fail(self, first, second):
     with self.assertRaises(AssertionError):
       asserts.assert_equal(first, second)
 
+class NotNoneAssertionsTest(parameterized.TestCase):
 
+    def test_assert_not_both_not_none(self):
+        asserts.assert_not_both_not_none(None, None)
+        asserts.assert_not_both_not_none(1, None)
+        asserts.assert_not_both_not_none(None, 1)
+
+        with self.assertRaises(AssertionError):
+            asserts.assert_not_both_not_none(1, 2)
 class IsDivisibleTest(parameterized.TestCase):
 
   def test_assert_is_divisible(self):