ソースを参照

add `csv.DictReader`

blueloveTH 2 年 前
コミット
f2e6004c1e
3 ファイル変更43 行追加2 行削除
  1. 9 1
      docs/modules/csv.md
  2. 24 1
      src/csv.cpp
  3. 10 0
      tests/80_csv.py

+ 9 - 1
docs/modules/csv.md

@@ -3,10 +3,13 @@ icon: package
 label: csv
 ---
 
-### `csv.reader(csvfile: list[str]) -> list`
+### `csv.reader(csvfile: list[str]) -> list[list]`
 
 Parse a CSV file into a list of lists.
 
+### `csv.DictReader(csvfile: list[str]) -> list[dict]`
+
+Parse a CSV file into a list of dictionaries.
 
 ## Example
 
@@ -22,4 +25,9 @@ print(csv.reader(data.splitlines()))
 #    ['a', 'b', 'c'],
 #    ['1', '2', '3']
 # ]
+
+print(csv.DictReader(data.splitlines()))
+# [
+#    {'a': '1', 'b': '2', 'c': '3'}
+# ]
 ```

+ 24 - 1
src/csv.cpp

@@ -6,7 +6,7 @@ namespace pkpy{
 void add_module_csv(VM *vm){
     PyObject* mod = vm->new_module("csv");
 
-    vm->bind(mod, "reader(csvfile: list[str]) -> list", [](VM* vm, ArgsView args){
+    vm->bind(mod, "reader(csvfile: list[str]) -> list[list]", [](VM* vm, ArgsView args){
         const List& csvfile = CAST(List&, args[0]);
         List ret;
         for(int i=0; i<csvfile.size(); i++){
@@ -57,6 +57,29 @@ void add_module_csv(VM *vm){
         }
         return VAR(std::move(ret));
     });
+
+    vm->bind(mod, "DictReader(csvfile: list[str]) -> list[dict]", [](VM* vm, ArgsView args){
+        PyObject* csv_reader = vm->_modules["csv"]->attr("reader");
+        PyObject* ret_obj = vm->call(csv_reader, args[0]);
+        const List& ret = CAST(List&, ret_obj);
+        if(ret.size() == 0){
+            vm->ValueError("empty csvfile");
+        }
+        List header = CAST(List&, ret[0]);
+        List new_ret;
+        for(int i=1; i<ret.size(); i++){
+            const List& row = CAST(List&, ret[i]);
+            if(row.size() != header.size()){
+                vm->ValueError("row.size() != header.size()");
+            }
+            Dict row_dict(vm);
+            for(int j=0; j<header.size(); j++){
+                row_dict.set(header[j], row[j]);
+            }
+            new_ret.push_back(VAR(std::move(row_dict)));
+        }
+        return VAR(std::move(new_ret));
+    });
 }
 
 }   // namespace pkpy

+ 10 - 0
tests/80_csv.py

@@ -26,3 +26,13 @@ test("""a,b,c,
 test("""a,b ,c,
 1,"22""33",3
 """,  [['a', 'b ', 'c', ''], ['1', '22"33', '3']])
+
+ret = csv.DictReader("""a,b,c
+1,2,3
+"4",5,6
+""".splitlines())
+
+assert list(ret)==[
+    {'a': '1', 'b': '2', 'c': '3'},
+    {'a': '4', 'b': '5', 'c': '6'},
+]