Merge pull request #8 from tuxxy/devel

Move 'generatemap.py' into 'getlocation.py'

getlocation.py

@@ -3,9 +3,26 @@
 from argparse import ArgumentParser
 from contextlib import contextmanager
 from pyipinfodb import pyipinfodb
+from mpl_toolkits.basemap import Basemap
+import matplotlib.pyplot as plt
 from sys import stdin, stdout, stderr
+import numpy as np
 
 
+def plot_locations(locations, output_filename):
+    ip_map = Basemap(projection='robin', lon_0=0, resolution='c')
+
+    for line in locations:
+    	srclong, srclat = map(float, line.split(','))
+    	x, y = ip_map(srclong, srclat)
+    	plt.plot(x,y, 'o', color='#ff0000', ms=2.7, markeredgewidth=0.5)
+
+
+    ip_map.drawcountries(color='#ffffff')
+    ip_map.fillcontinents(color='#cccccc',lake_color='#ffffff')
+
+    plt.savefig(output_filename, dpi=600)
+
 def memoize(func):
     '''
     Decorator for a function which caches results. Based on
@@ -43,15 +60,20 @@ def main():
     parser = ArgumentParser()
     parser.add_argument('--input', '-i', default=stdin)
     parser.add_argument('--output', '-o', default=stdout)
+    parser.add_argument('--gen-map', '-G', action='store_true')
     parser.add_argument('API_KEY')
     args = parser.parse_args()
 
     ip_lookup = pyipinfodb.IPInfo(args.API_KEY)
 
     with smart_open(args.input) as istr:
         with smart_open(args.output, 'w') as ostr:
-            for ip in istr:
-                print(*get_location(ip_lookup, ip), sep=",", file=ostr)
+            if args.gen_map:
+                for ip in istr:
+                    plot_locations(*(get_location(ip_lookup, ip), "ip_map.png"))
+            else:
+                for ip in istr:
+                    print(*get_location(ip_lookup, ip), sep=",", file=ostr)
 
 if __name__ == '__main__':
     main()