sim.pysim: in write_vcd(), close files if an exception is raised.

[nmigen.git] / nmigen / sim / pysim.py

diff --git a/nmigen/sim/pysim.py b/nmigen/sim/pysim.py
index 5f30816767e51066f36033d40f073173353b6105..847a5a1e6c0229f1ad48803808e1aa4262403b9d 100644 (file)
--- a/nmigen/sim/pysim.py
+++ b/nmigen/sim/pysim.py
@@ -402,10 +402,16 @@ class Simulator:
             Signals to display traces for.
         """
         if self._state.timeline.now != 0.0:
+            for file in (vcd_file, gtkw_file):
+                if hasattr(file, "close"):
+                    file.close()
             raise ValueError("Cannot start writing waveforms after advancing simulation time")
+
         waveform_writer = _VCDWaveformWriter(self._fragment,
             vcd_file=vcd_file, gtkw_file=gtkw_file, traces=traces)
-        self._waveform_writers.append(waveform_writer)
-        yield
-        waveform_writer.close(self._state.timeline.now)
-        self._waveform_writers.remove(waveform_writer)
+        try:
+            self._waveform_writers.append(waveform_writer)
+            yield
+        finally:
+            waveform_writer.close(self._state.timeline.now)
+            self._waveform_writers.remove(waveform_writer)