add csv.DictReader

This commit is contained in:
blueloveTH 2023-11-29 00:51:02 +08:00
parent 657062d25a
commit f2e6004c1e
3 changed files with 43 additions and 2 deletions

View File

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

View File

@ -6,7 +6,7 @@ namespace pkpy{
void add_module_csv(VM *vm){ void add_module_csv(VM *vm){
PyObject* mod = vm->new_module("csv"); 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]); const List& csvfile = CAST(List&, args[0]);
List ret; List ret;
for(int i=0; i<csvfile.size(); i++){ for(int i=0; i<csvfile.size(); i++){
@ -57,6 +57,29 @@ void add_module_csv(VM *vm){
} }
return VAR(std::move(ret)); 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 } // namespace pkpy

View File

@ -26,3 +26,13 @@ test("""a,b,c,
test("""a,b ,c, test("""a,b ,c,
1,"22""33",3 1,"22""33",3
""", [['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'},
]