From f2c1e279c96b8cfa79ee015ff0c5b125ef07b9a3 Mon Sep 17 00:00:00 2001
From: Ivan Kravets <me@ikravets.com>
Date: Sat, 11 Jun 2016 00:55:38 +0300
Subject: [PATCH] Better checking of program size before uploading // Issue
 #689

---
 platformio/builder/main.py             | 3 ---
 platformio/builder/scripts/atmelavr.py | 5 +++--
 platformio/builder/scripts/atmelsam.py | 5 +++--
 platformio/builder/tools/pioupload.py  | 5 ++---
 4 files changed, 8 insertions(+), 10 deletions(-)

diff --git a/platformio/builder/main.py b/platformio/builder/main.py
index d060aab2..8b01c7b2 100644
--- a/platformio/builder/main.py
+++ b/platformio/builder/main.py
@@ -171,6 +171,3 @@ if "envdump" in COMMAND_LINE_TARGETS:
 if "idedata" in COMMAND_LINE_TARGETS:
     print json.dumps(env.DumpIDEData())
     env.Exit()
-
-if set(["upload", "uploadlazy", "program"]) & set(COMMAND_LINE_TARGETS):
-    env.CheckUploadSize()
diff --git a/platformio/builder/scripts/atmelavr.py b/platformio/builder/scripts/atmelavr.py
index ae84cf21..97085a03 100644
--- a/platformio/builder/scripts/atmelavr.py
+++ b/platformio/builder/scripts/atmelavr.py
@@ -142,7 +142,7 @@ AlwaysBuild(target_size)
 #
 
 upload = env.Alias(["upload", "uploadlazy"], target_firm,
-                   [BeforeUpload, "$UPLOADHEXCMD"])
+                   [env.CheckUploadSize, BeforeUpload, "$UPLOADHEXCMD"])
 AlwaysBuild(upload)
 
 #
@@ -159,7 +159,8 @@ AlwaysBuild(uploadeep)
 # Target: Upload firmware using external programmer
 #
 
-program = env.Alias("program", target_firm, [BeforeUpload, "$PROGRAMHEXCMD"])
+program = env.Alias("program", target_firm,
+                    [env.CheckUploadSize, BeforeUpload, "$PROGRAMHEXCMD"])
 AlwaysBuild(program)
 
 #
diff --git a/platformio/builder/scripts/atmelsam.py b/platformio/builder/scripts/atmelsam.py
index 79eca59e..f5ad7866 100644
--- a/platformio/builder/scripts/atmelsam.py
+++ b/platformio/builder/scripts/atmelsam.py
@@ -174,10 +174,11 @@ AlwaysBuild(target_size)
 #
 
 if env.subst("$BOARD") == "zero":
-    upload = env.Alias(["upload", "uploadlazy"], target_firm, "$UPLOADCMD")
+    upload = env.Alias(["upload", "uploadlazy"], target_firm,
+                       [env.CheckUploadSize, "$UPLOADCMD"])
 else:
     upload = env.Alias(["upload", "uploadlazy"], target_firm,
-                       [BeforeUpload, "$UPLOADCMD"])
+                       [env.CheckUploadSize, BeforeUpload, "$UPLOADCMD"])
 
 AlwaysBuild(upload)
 
diff --git a/platformio/builder/tools/pioupload.py b/platformio/builder/tools/pioupload.py
index b75501c5..12b083a5 100644
--- a/platformio/builder/tools/pioupload.py
+++ b/platformio/builder/tools/pioupload.py
@@ -133,7 +133,7 @@ def UploadToDisk(_, target, source, env):  # pylint: disable=W0613,W0621
           "Please restart your board.")
 
 
-def CheckUploadSize(env):
+def CheckUploadSize(_, target, source, env):  # pylint: disable=W0613,W0621
     max_size = int(env.get("BOARD_OPTIONS", {}).get("upload", {}).get(
         "maximum_size", 0))
     if max_size == 0 or "SIZETOOL" not in env:
@@ -141,8 +141,7 @@ def CheckUploadSize(env):
 
     sysenv = environ.copy()
     sysenv['PATH'] = str(env['ENV']['PATH'])
-    cmd = [env.subst("$SIZETOOL"), "-B"]
-    cmd.append(env.subst(join("$BUILD_DIR", "$PROGNAME$PROGSUFFIX")))
+    cmd = [env.subst("$SIZETOOL"), "-B", str(source[0])]
     result = util.exec_command(cmd, env=sysenv)
     if result['returncode'] != 0:
         return