From a941071c926542562968bac627db9a44999e53b0 Mon Sep 17 00:00:00 2001
From: Mike Naberezny <mike@naberezny.com>
Date: Sun, 14 Dec 2014 16:44:13 -0800
Subject: [PATCH] Handle overflow when setting registers

---
 py65/monitor.py            | 24 +++++++++++++++++-------
 py65/tests/test_monitor.py | 16 ++++++++++++++++
 2 files changed, 33 insertions(+), 7 deletions(-)

diff --git a/py65/monitor.py b/py65/monitor.py
index 5f9a886..96a0687 100644
--- a/py65/monitor.py
+++ b/py65/monitor.py
@@ -223,8 +223,7 @@ class Monitor(cmd.Cmd):
         self._output("\n" + repr(self._mpu))
 
     def _output(self, stuff):
-        if stuff is not None:
-            self.stdout.write(stuff + "\n")
+        self.stdout.write("%s\n" % stuff)
 
     def _exit(self, exitcode=0):
         sys.exit(exitcode)
@@ -523,12 +522,23 @@ class Monitor(cmd.Cmd):
                 self._output("Invalid register: %s" % register)
             else:
                 try:
-                    intval = self._address_parser.number(value) & self.addrMask
-                    if len(register) == 1:
-                        intval &= self.byteMask
-                    setattr(self._mpu, register, intval)
-                except KeyError as exc:
+                    intval = self._address_parser.number(value)
+                except KeyError as exc: # label not found
                     self._output(exc.args[0])
+                    continue
+                except OverflowError as exc: # wider than address space
+                    overflow = True
+                    msg = "Overflow: %r too wide for register %r"
+                    self._output(msg % (value, register))
+                    continue
+
+                if register != 'pc':
+                    if intval != (intval & self.byteMask):
+                        msg = "Overflow: %r too wide for register %r"
+                        self._output(msg % (value, register))
+                        continue
+
+                setattr(self._mpu, register, intval)
 
     def help_cd(self):
         self._output("cd <directory>")
diff --git a/py65/tests/test_monitor.py b/py65/tests/test_monitor.py
index eb3bd27..7365e9e 100644
--- a/py65/tests/test_monitor.py
+++ b/py65/tests/test_monitor.py
@@ -854,6 +854,22 @@ class MonitorTests(unittest.TestCase):
         self.assertEqual(0x46, mon._mpu.sp)
         self.assertEqual(0x4600, mon._mpu.pc)
 
+    def test_registers_pc_overflow(self):
+        stdout = StringIO()
+        mon = Monitor(stdout=stdout)
+        mon.do_registers('pc=10000')
+        out = stdout.getvalue()
+        expected = "Overflow: '10000' too wide for register 'pc'"
+        self.assertTrue(out.startswith(expected))
+
+    def test_registers_a_overflow(self):
+        stdout = StringIO()
+        mon = Monitor(stdout=stdout)
+        mon.do_registers('a=100')
+        out = stdout.getvalue()
+        expected = "Overflow: '100' too wide for register 'a'"
+        self.assertTrue(out.startswith(expected))
+
     def test_help_registers(self):
         stdout = StringIO()
         mon = Monitor(stdout=stdout)