amaranth-lang · Jan 31, 2020
Showing with 43 additions and 4 deletions.

+26 −4 nmigen/hdl/dsl.py

+17 −0 nmigen/test/test_hdl_dsl.py
diff --git a/nmigen/hdl/dsl.py b/nmigen/hdl/dsl.py
@@ -1,6 +1,7 @@
 from collections import OrderedDict, namedtuple
 from collections.abc import Iterable
-from contextlib import contextmanager
+from contextlib import contextmanager, _GeneratorContextManager
+from functools import wraps
 from enum import Enum
 import warnings
 
@@ -113,6 +114,27 @@ def __setattr__(self, name, domain):
         self._builder._add_domain(domain)
 
 
+# It's not particularly clean to depend on an internal interface, but, unfortunately, __bool__
+# must be defined on a class to be called during implicit conversion.
+class _GuardedContextManager(_GeneratorContextManager):
+    def __init__(self, keyword, func, args, kwds):
+        self.keyword = keyword
+        return super().__init__(func, args, kwds)
+
+    def __bool__(self):
+        raise SyntaxError("`if m.{kw}(...):` does not work; use `with m.{kw}(...)`"
+                          .format(kw=self.keyword))
+
+
+def _guardedcontextmanager(keyword):
+    def decorator(func):
+        @wraps(func)
+        def helper(*args, **kwds):
+            return _GuardedContextManager(keyword, func, args, kwds)
+        return helper
+    return decorator
+
+
 class FSM:
     def __init__(self, state, encoding, decoding):
         self.state    = state
@@ -183,7 +205,7 @@ def _check_signed_cond(self, cond):
                           SyntaxWarning, stacklevel=4)
         return cond
 
-    @contextmanager
+    @_guardedcontextmanager("If")
     def If(self, cond):
         self._check_context("If", context=None)
         cond = self._check_signed_cond(cond)
@@ -206,7 +228,7 @@ def If(self, cond):
             self.domain._depth -= 1
             self._statements = _outer_case
 
-    @contextmanager
+    @_guardedcontextmanager("Elif")
     def Elif(self, cond):
         self._check_context("Elif", context=None)
         cond = self._check_signed_cond(cond)
@@ -226,7 +248,7 @@ def Elif(self, cond):
             self.domain._depth -= 1
             self._statements = _outer_case
 
-    @contextmanager
+    @_guardedcontextmanager("Else")
     def Else(self):
         self._check_context("Else", context=None)
         src_loc = tracer.get_src_loc(src_loc_at=1)

diff --git a/nmigen/test/test_hdl_dsl.py b/nmigen/test/test_hdl_dsl.py
@@ -300,6 +300,23 @@ def test_Elif_signed_suspicious(self):
             with m.Elif(~True):
                 pass
 
+    def test_if_If_Elif_Else(self):
+        m = Module()
+        with self.assertRaises(SyntaxError,
+                msg="`if m.If(...):` does not work; use `with m.If(...)`"):
+            if m.If(0):
+                pass
+        with m.If(0):
+            pass
+        with self.assertRaises(SyntaxError,
+                msg="`if m.Elif(...):` does not work; use `with m.Elif(...)`"):
+            if m.Elif(0):
+                pass
+        with self.assertRaises(SyntaxError,
+                msg="`if m.Else(...):` does not work; use `with m.Else(...)`"):
+            if m.Else():
+                pass
+
     def test_Switch(self):
         m = Module()
         with m.Switch(self.w1):