Fixed ant_landscape/test.py so that it can be imported

This allows the `bl_load_py_modules.py` unit test to do its job.

1 files changed, 17 insertions, 15 deletions

diff --git a/ant_landscape/test.py b/ant_landscape/test.py
index 32894a18..d56815f2 100644
--- a/ant_landscape/test.py
+++ b/ant_landscape/test.py
@@ -1,19 +1,21 @@
-from stats import Stats
 from numpy import *
 
-stats = Stats()
-
-a = zeros(10000000)
-print(stats.time())
-print(stats.memory())
-a = sin(a)
-print(stats.time())
-print(stats.memory())
-a = cos(a)
-print(stats.time())
-print(stats.memory())
-a = cos(a)**2+sin(a)**2
-print(stats.time())
-print(stats.memory())
+if __name__ == '__main__':
+    # This only works when this script is loaded as main, or
+    # run directly from the ant_landscape directory.
+    from stats import Stats
 
+    stats = Stats()
 
+    a = zeros(10000000)
+    print(stats.time())
+    print(stats.memory())
+    a = sin(a)
+    print(stats.time())
+    print(stats.memory())
+    a = cos(a)
+    print(stats.time())
+    print(stats.memory())
+    a = cos(a)**2+sin(a)**2
+    print(stats.time())
+    print(stats.memory())