Merge pull request #21 from zuzhi/master

支持指定经纬度字段

Author: wandergis

README.md

@@ -6,7 +6,7 @@
 
 需要js版本可以移步[coordtransform](https://github.com/wandergis/coordtransform)
 
-## 使用说明
+## 使用说明(coordTransform_utils.py)
 
 ### 方法说明
 
@@ -27,22 +27,46 @@ g.geocode('北京市朝阳区朝阳公园')
 ### 测试
 
 ```bash
-# 测试
+# 测试(转换坐标 128.543,37.065 )
 $ python coordTransform_utils.py
 [128.54944656269413, 37.07113427883019] [128.5365893261212, 37.058754503281534] [128.54820547949757, 37.065651049489816] [128.53779452050244, 37.06434895051018] [128.53136876750008, 37.0580926428705] [128.55468192918485, 37.07168344938498] None
 ```
 
-## 批量转换csv文件(coord_converter.py)
+## 批量转换csv文件使用说明(coord_converter.py)
 
-使用方法：
+### 使用说明：
 
 ```bash
+# 查看使用帮助
 $ python coord_converter.py -h
 
-Usage: coord_converter.py -i <input> -o <output> -t <type>
+usage: coord_converter.py [-h] -i INPUT -o OUTPUT -t TYPE [-n LNG_COLUMN] [-a LAT_COLUMN] [-s SKIP_INVALID_ROW]
 
-where <type> is one of:
-    g2b, b2g, w2g, g2w, b2w, w2b
+Convert coordinates in csv files.
 
-Example: coord_converter.py -i /path/to/input_file.csv -o /path/to/output_file.csv -t b2g
+optional arguments:
+  -h, --help            show this help message and exit
+
+arguments:
+  -i , --input          Location of input file
+  -o , --output         Location of output file
+  -t , --type           Convert type, must be one of: g2b, b2g, w2g, g2w, b2w,
+                        w2b
+  -n , --lng_column     Column name for longitude (default: lng)
+  -a , --lat_column     Column name for latitude (default: lat)
+  -s , --skip_invalid_row
+                        Whether to skip invalid row (default: False)
+```
+
+### 示例
+
+```bash
+# 不指定经纬度列名（默认为'lng', 'lat'）
+$ python coord_converter.py -i test_input.csv -o test_output.csv -t b2g
+
+# 指定经纬度列名
+$ python coord_converter.py -i test_input.csv -o test_output.csv -t b2g -n 经度 -a 纬度
+
+# 跳过无效经纬度的行（默认不跳过）
+$ python coord_converter.py -i test_input.csv -o test_output.csv -t b2g -n 经度 -a 纬度 -s True
 ```

coord_converter.py

@@ -1,35 +1,80 @@
 #!/usr/bin/python
+# -*- coding: utf-8 -*-
 
 import csv
 import sys
-import getopt
+import argparse
 from coordTransform_utils import gcj02_to_bd09
 from coordTransform_utils import bd09_to_gcj02
 from coordTransform_utils import wgs84_to_gcj02
 from coordTransform_utils import gcj02_to_wgs84
 from coordTransform_utils import bd09_to_wgs84
 from coordTransform_utils import wgs84_to_bd09
 
-# Configuation
+# Configuration
 # Input file name
 INPUT = ''
 # Output file name
 OUTPUT = ''
 # Convert type: g2b, b2g, w2g, g2w, b2w, w2b
 TYPE = ''
+# lng column name
+LNG_COLUMN = ''
+# lat column name
+LAT_COLUMN = ''
+# Skip invalid row
+SKIP_INVALID_ROW = False
 
 def convert():
     with open(INPUT, 'r') as input_file:
         input_file_reader = csv.reader(input_file)
+        headers = next(input_file_reader)
+        lng_index, lat_index = get_lng_lat_index(headers)
         results = []
-        for lat, lng in input_file_reader:
-            result = convert_by_type(float(lng), float(lat), TYPE)
+
+        for index, row in enumerate(input_file_reader):
+            result = []
+            try:
+                result = convert_by_type(float(row[lng_index]), float(row[lat_index]), TYPE)
+            except ValueError:
+                # Deal with ValueError(invalid lng or lat)
+                # print(index + 2, row[lng_index], row[lat_index]) # '+ 2' is due to zero-based index and first row is header
+                result = row[lng_index], row[lat_index]
             results.append(result)
 
     with open(OUTPUT, 'w') as output_file:
         output_file_writer = csv.writer(output_file)
-        for result in results:
-            output_file_writer.writerow(result)
+
+        with open(INPUT, 'r') as input_file:
+            input_file_reader = csv.reader(input_file)
+            headers = next(input_file_reader)
+            lng_index, lat_index = get_lng_lat_index(headers)
+
+            output_file_writer.writerow(headers)
+            for index, row in enumerate(input_file_reader):
+                row[lng_index] = results[index][0]
+                row[lat_index] = results[index][1]
+                if type(row[lng_index]) is not float or type(row[lat_index]) is not float:
+                    # Data is invalid
+                    if SKIP_INVALID_ROW:
+                        # Skip invalid row
+                        pass
+                    else:
+                        # Reserve invalid row
+                        output_file_writer.writerow(row)
+                else:
+                    # Data is valid
+                    output_file_writer.writerow(row)
+
+def get_lng_lat_index(headers):
+    try:
+        if LNG_COLUMN == '' and LAT_COLUMN == '':
+            return [headers.index('lng'), headers.index('lat')]
+        else:
+            return [headers.index(LNG_COLUMN), headers.index(LAT_COLUMN)]
+    except ValueError as error:
+        print('Error: ' + str(error).split('is', 1)[0] + 'is missing from csv header. Or use -n or -a to specify custom column name for lng or lat.')
+        sys.exit()
 
 def convert_by_type(lng, lat, type):
     if type == 'g2b':
@@ -45,37 +90,38 @@ def convert_by_type(lng, lat, type):
     elif type == 'w2b':
         return wgs84_to_bd09(lng, lat)
     else:
-        print('Usage: The <type> must be in one of g2b, b2g, w2g, g2w, b2w, w2b')
+        print('Usage: type must be in one of g2b, b2g, w2g, g2w, b2w, w2b')
         sys.exit()
 
-def usage():
-    print('Usage: coord_converter.py -i <input> -o <output> -t <type>')
-    print('')
-    print('where <type> is one of:')
-    print('    g2b, b2g, w2g, g2w, b2w, w2b')
-    print('')
-    print('Example: coord_converter.py -i /path/to/input_file.csv -o /path/to/output_file.csv -t b2g')
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Convert coordinates in csv files.', usage='%(prog)s [-h] -i INPUT -o OUTPUT -t TYPE [-n LNG_COLUMN] [-a LAT_COLUMN] [-s SKIP_INVALID_ROW]', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
-if __name__ == "__main__":
-    argv = sys.argv[1:]
-    try:
-        opts, args = getopt.getopt(argv, "hi:o:t:", ["input=", "output=", "type="])
-    except getopt.GetoptError:
-        usage()
-        sys.exit(2)
-    for opt, arg in opts:
-        if opt == '-h':
-            usage()
-            sys.exit()
-        elif opt in ("-i", "--input"):
-            INPUT = arg
-        elif opt in ("-o", "--output"):
-            OUTPUT = arg
-        elif opt in ("-t", "--type"):
-            TYPE = arg
-
-    if not (INPUT and OUTPUT and TYPE):
-        usage()
-        sys.exit()
+    group = parser.add_argument_group('arguments')
+
+    group.add_argument('-i', '--input', help='Location of input file', default=argparse.SUPPRESS, metavar='')
+    group.add_argument('-o', '--output', help='Location of output file', default=argparse.SUPPRESS, metavar='')
+    group.add_argument('-t', '--type', help='Convert type, must be one of: g2b, b2g, w2g, g2w, b2w, w2b', default=argparse.SUPPRESS, metavar='')
+    group.add_argument('-n', '--lng_column', help='Column name for longitude', default='lng', metavar='')
+    group.add_argument('-a', '--lat_column', help='Column name for latitude', default='lat', metavar='')
+    group.add_argument('-s', '--skip_invalid_row', help='Whether to skip invalid row', default=False, type=bool, metavar='')
+
+    args = parser.parse_args()
+    # print('\nArguments you provide are:')
+    # for arg in vars(args):
+    #     print '{0:20} {1}'.format(arg, str(getattr(args, arg)))
+
+    # Get arguments
+    if not args.input or not args.output or not args.type:
+        parser.print_help()
+    else:
+        INPUT = args.input
+        OUTPUT = args.output
+        TYPE = args.type
+
+    if args.lng_column and args.lat_column:
+        LNG_COLUMN, LAT_COLUMN = args.lng_column, args.lat_column
+
+    if args.skip_invalid_row:
+        SKIP_INVALID_ROW = args.skip_invalid_row
 
     convert()

test_input.csv

@@ -0,0 +1,3 @@
+lng,lat,column1
+116.4172,39.93889,value1
+102.72569,25.0444,value2