py-02-csv-groupby
1.000
Challenge · difficulty 2/5
# CSV Group Sum
Implement a file **`solution.py`** containing a function:
```python
def group_sum(csv_text: str, key_col: str, val_col: str) -> dict[str, float]:
...
```
Parse `csv_text` as CSV and return a dictionary mapping each distinct value in
column `key_col` to the **sum** of column `val_col` for all rows with that key.
Rules:
- The **first non-blank line** is the header row naming the columns.
- Columns are comma-separated. `key_col` and `val_col` name two of those columns.
- Values in `val_col` are numeric (int or float); sum them as floats.
- **Ignore blank lines** anywhere in the input (including trailing newlines).
- If the input has only a header (or is empty), return an empty dict `{}`.
- You may assume `key_col` and `val_col` exist in the header.
Example:
```python
csv_text = "name,amount\\nalice,10\\nbob,5\\nalice,2.5\\n"
group_sum(csv_text, "name", "amount") == {"alice": 12.5, "bob": 5.0}
```
tests/test_group_sum.py
from solution import group_sum
def test_basic_sum():
csv_text = "name,amount\nalice,10\nbob,5\nalice,2.5\n"
assert group_sum(csv_text, "name", "amount") == {"alice": 12.5, "bob": 5.0}
def test_returns_floats():
csv_text = "k,v\na,1\na,2\n"
out = group_sum(csv_text, "k", "v")
assert out == {"a": 3.0}
assert all(isinstance(v, float) for v in out.values())
def test_ignores_blank_lines():
csv_text = "\n\nname,amount\nalice,1\n\nbob,2\n\n\n"
assert group_sum(csv_text, "name", "amount") == {"alice": 1.0, "bob": 2.0}
def test_header_only():
assert group_sum("name,amount\n", "name", "amount") == {}
def test_empty_input():
assert group_sum("", "name", "amount") == {}
assert group_sum("\n\n", "name", "amount") == {}
def test_column_order_independent():
csv_text = "amount,name\n10,x\n5,y\n3,x\n"
assert group_sum(csv_text, "name", "amount") == {"x": 13.0, "y": 5.0}
def test_three_columns_picks_right_ones():
csv_text = "region,product,sales\neast,a,100\nwest,b,50\neast,c,25\n"
assert group_sum(csv_text, "region", "sales") == {"east": 125.0, "west": 50.0}
def test_negative_and_float_values():
csv_text = "g,n\nx,-1.5\nx,0.5\ny,4\n"
assert group_sum(csv_text, "g", "n") == {"x": -1.0, "y": 4.0}
System prompt
You are an expert programmer. Solve the task exactly as specified. Output your solution as fenced code blocks using the required file name(s) and the exact function/type signatures requested. Prefer correctness; do not include prose outside code unless asked.Proposed solution
```python
import csv
from io import StringIO
def group_sum(csv_text: str, key_col: str, val_col: str) -> dict[str, float]:
reader = csv.reader(StringIO(csv_text))
# Find header (first non-blank row)
header = None
for row in reader:
if row:
header = row
break
if header is None:
return {}
key_idx = header.index(key_col)
val_idx = header.index(val_col)
result = {}
for row in reader:
if not row:
continue
key = row[key_idx]
value = float(row[val_idx])
result[key] = result.get(key, 0.0) + value
return result
```Test output (stdout)
........ [100%]
8 passed in 0.01s
Reasoning (chain-of-thought)
The user wants a Python function `group_sum` that:
1. Takes CSV text as a string, a key column name, and a value column name.
2. Parses the CSV.
3. Returns a dictionary mapping each distinct key value to the sum of the corresponding value column.
Key requirements:
- First non-blank line is the header.
- Comma-separated.
- Values in val_col are numeric, sum as floats.
- Ignore blank lines.
- If only header or empty, return {}.
- Assume key_col and val_col exist in header.
I'll use Python's `csv` module for robust parsing.
Steps:
1. Split the text into lines.
2. Filter out blank lines.
3. The first non-blank line is the header. Parse it to get column names.
4. Find indices of key_col and val_col.
5. Iterate through remaining lines, parse each row, accumulate sums.
6. Return the dictionary.
Let me write the code:
```python
import csv
from io import StringIO
def group_sum(csv_text: str, key_col: str, val_col: str) -> dict[str, float]:
lines = csv_text.strip().split('\n')
# Filter out blank lines
non_blank_lines = [line for line in lines if line.strip()]
if not non_blank_lines:
return {}
# First non-blank line is header
header_line = non_blank_lines[0]
reader = csv.reader(StringIO(header_line))
header = next(reader)
key_idx = header.index(key_col)
val_idx = header.index(val_col)
result = {}
for line in non_blank_lines[1:]:
reader = csv.reader(StringIO(line))
row = next(reader)
key = row[key_idx]
value = float(row[val_idx])
result[key] = result.get(key, 0.0) + value
return result
```
Wait, I should be more careful with CSV parsing. Using `csv.reader` on the entire text might be better to handle edge cases like quoted fields with newlines. But the problem says "first non-blank line is the header" and "ignore blank lines", which suggests a simpler line-based approach might be expected. However, using `csv.reader` on the whole text is more robust.
Let me reconsider: if I use `csv.reader` on the whole text, blank lines would produce empty rows `[]`. I can filter those out.
Actually, let me use a simpler approach that's more straightforward:
```python
import csv
from io import StringIO
def group_sum(csv_text: str, key_col: str, val_col: str) -> dict[str, float]:
reader = csv.reader(StringIO(csv_text))
# Find header (first non-blank row)
header = None
for row in reader:
if row: # non-blank
header = row
break
if header is None:
return {}
key_idx = header.index(key_col)
val_idx = header.index(val_col)
result = {}
for row in reader:
if not row: # skip blank lines
continue
key = row[key_idx]
value = float(row[val_idx])
result[key] = result.get(key, 0.0) + value
return result
```
This should work correctly. Let me verify with the example:
- csv_text = "name,amount\nalice,10\nbob,5\nalice,2.5\n"
- Header: ["name", "amount"]
- key_idx = 0, val_idx = 1
- Row ["alice", "10"]: result["alice"] = 10.0
- Row ["bob", "5"]: result["bob"] = 5.0
- Row ["alice", "2.5"]: result["alice"] = 12.5
- Return {"alice": 12.5, "bob": 5.0} ✓
This looks correct. I'll output this solution.