Repository: koaning/clumper

Issue:
Data Writer: json/jsonl
Data Loader: .to_json()/to_jsonl()

It'd be nice if we could also write data to disk. A syntax like this would be nice: 

```python
Clumper.to_json(path, settings)
Clumper.to_jsonl(path, settings)
```

Discussion:

Patch:
diff --git a/clumper/clump.py b/clumper/clump.py
index 4ea124c..fe729da 100644
--- a/clumper/clump.py
+++ b/clumper/clump.py
@@ -6,6 +6,7 @@ import urllib.request
 from functools import reduce
 from statistics import mean, variance, stdev, median
 
+
 from clumper.decorators import return_value_if_empty, grouped, dict_collection_only
 
 
@@ -130,8 +131,68 @@ class Clumper:
         except Exception:
             raise RuntimeError("Error occured during reading in JSONL file")
 
+    def write_json(self, path, sort_keys=False, indent=None):
+        """
+        Writes to a json file.
+
+        Arguments:
+            path: filename
+            sort_keys: If sort_keys is true (default: False), then the output of dictionaries will be sorted by key.
+            indent: If indent is a non-negative integer (default: None), then JSON array elements members will be pretty-printed with that indent level.
+
+        Usage:
+
+        ```python
+        from clumper import Clumper
+        clump_orig = Clumper.read_json("tests/data/pokemon.json")
+        clump_orig.write_jsonl("tests/data/pokemon_copy.json")
+
+        clump_copy = Clumper.read_json("tests/data/pokemon_copy.json")
+        assert clump_copy.collect() == clump_orig.collect()
+        """
+
+        try:
+            # Create a new file and open it for writing
+            with open(path, "w") as f:
+                json.dump(self.collect(), f, sort_keys=sort_keys, indent=indent)
+        except Exception:
+            raise RuntimeError("Error occured during writing JSON file")
+
+    def write_jsonl(self, path, sort_keys=False, indent=None):
+        """
+        Writes to a jsonl file.
+
+        Arguments:
+            path: filename
+            sort_keys: If sort_keys is true (default: False), then the output of dictionaries will be sorted by key.
+            indent: If indent is a non-negative integer (default: None), then JSON array elements members will be pretty-printed with that indent level.
+        Usage:
+
+        ```python
+        from clumper import Clumper
+        clump_orig = Clumper.read_jsonl("tests/data/cards.jsonl")
+        clump_orig.write_jsonl("tests/data/cards_copy.jsonl")
+
+        clump_copy = Clumper.read_jsonl("tests/data/cards_copy.jsonl")
+
+        assert clump_copy.collect() == clump_orig.collect()
+        """
+
+        try:
+            # Create a new file and open it for writing
+            with open(path, "x") as f:
+                for current_line_nr, json_dict in enumerate(self.collect()):
+                    f.write(
+                        json.dumps(json_dict, sort_keys=sort_keys, indent=indent) + "\n"
+                    )
+
+        except Exception:
+            raise RuntimeError("Error occured during writing JSONL file")
+
     @classmethod
-    def read_csv(cls, path, delimiter=",", fieldnames=None, n=None):
+    def read_csv(
+        cls, path, delimiter=",", na_values=None, dtype=None, fieldnames=None, n=None
+    ):
         """
         Reads in a csv file. Can also read files from url.
 
@@ -145,8 +206,19 @@ class Clumper:
                         row of the csv will provide the Clumper keys if fieldnames is `None`. If fieldnames
                         is provided, then the first row becomes part of the data. You should ensure that
                         the correct number of fieldnames is supplied, as an incorrect number can lead
-                        to truncation. If you have seven columns and your fieldnames length is 3,
-                        then every row will have only 3 values, the remaining four will be cut off.
+                        to an irregular outcome. If the row has seven fields and the number of fields in
+                        fieldnames length is 3, then every row will have only 3 values, the remaining four
+                        will be lumped into a list, and assigned key `None`. If the rows have fewer fields
+                        than fieldnames, then the missing values are filled in with `None`.
+            na_values:  This provides an option for treating null values. If `ignore`, null values are
+                        returned as empty strings (""). If `None`, then for each row, the key,value pair
+                        with the null values  will be truncated from the row. The only values treated as
+                        null are empty strings("") and "NA".
+            dtype: Data type for each value in a key:value pair. If `None`, then values will be read in as strings.
+                   Available dtypes are (int, float, str). If a single dtype is passed, then all values will be
+                   converted to the data type and raise an error, if not applicable. For different data types for different
+                   key, value pairs, a dictionary of {key: data_type} passed to dtype argument will change the value for
+                   every key with the data type, and raise an error if not applicable.
 
         Usage:
 
@@ -174,19 +246,102 @@ class Clumper:
             if n <= 0:
                 raise ValueError("Number of lines to read must be > 0.")
 
-        if path.startswith("https:") or path.startswith("http:"):
+        # conveniently excludes pathlib files here and removes
+        # the need to write code to check pathlib files in other places.
+        if isinstance(path, str) and path.startswith(("https:", "http:")):
             with urllib.request.urlopen(path) as resp:
                 if fieldnames is None:
                     fieldnames = resp.readline().decode().strip().split(",")
-                # this section allows us to chunk the rows, if nrows is supplied
+                # This section allows us to chunk the rows, if nrows is supplied.
                 body = it.islice(resp, 0, n)
                 body = (word.decode().strip().split(",") for word in body)
                 body = it.product([fieldnames], body)
-                return Clumper([dict(zip(key, values)) for key, values in body])
+                result = [dict(zip(key, values)) for key, values in body]
+        else:
+            with open(path, newline="") as csvfile:
+                reader = csv.DictReader(
+                    csvfile, delimiter=delimiter, fieldnames=fieldnames
+                )
+                # python version less than 3.8 returns an OrderedDict
+                result = [dict(entry) for entry in it.islice(reader, 0, n)]
+
+        # Null values, same as missing keys.
+        # If there are null values/missing keys, they will be truncated from the dictionary.
+        # Python's csv module treats null values as empty strings when writing to a csv -
+        # https://docs.python.org/3.8/library/csv.html#csv.DictWriter.
+        # The user can choose to explicitly show missing keys/null values in the dictionary,
+        # by assigning `ignore` to the na_values argument. At the moment, the default for
+        # null values are empty string ("") and "NA".
+
+        if na_values == "ignore":
+            result = result
+        else:
+            na_values = ["", "NA"]
+            result = [
+                {key: value for key, value in entry.items() if value not in na_values}
+                for entry in result
+            ]
+
+        # The csv module reads every row as a string, there are no data type assumptions.
+        # This function attempts to solve this. The user can pass a string of either
+        # ('int', 'str', 'float') or if the user knows the keys/fieldnames, can pass a
+        # dictionary mapping the key to the data type.
+        # Technically 'str' data type is not needed, since data is read in as strings anyway.
+
+        if not (isinstance(dtype, (dict, str)) or dtype is None):
+            raise TypeError(
+                """dtype should be a dictionary pair of key and data type, or a single string data type"""
+            )
+
+        dtype_mapping = {"int": int, "float": float, "str": str}
+        if dtype:
+            if isinstance(dtype, str) and dtype in ("int", "float", "str"):
+                result = [
+                    {key: dtype_mapping[dtype](value) for key, value in entry.items()}
+                    for entry in result
+                ]
+            else:
+                result = [
+                    {
+                        key: dtype_mapping[dtype[key]](value) if key in dtype else value
+                        for key, value in entry.items()
+                    }
+                    for entry in result
+                ]
+
+        return Clumper(result)
+
+    def write_csv(self, path, mode="w"):
+        """
+        Write to a csv file.
+
+        Arguments:
+        path: filename
+        mode: `w` writes to a file if it does not exist, or overwrites if it already exists,
+               while `a`: - append to file if it already exists. The default is `w`.
+
+        Note that null values will be exported as empty strings; this is the convention chosen by Python.
+
+        Usage:
+
+        ```python
+        from clumper import Clumper
+        from pathlib import Path
+        path = '/tmp/monopoly.csv'
+        Clumper.read_csv("tests/data/monopoly.csv").write_csv(path)
+        reader = Clumper.read_csv(path)
+        assert Clumper.read_csv("tests/data/monopoly.csv").collect() == reader.collect()
+        ```
+        """
+
+        with open(path, mode=mode, newline="") as csvfile:
+            fieldnames = self.keys()
+            writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+
+            writer.writeheader()
 
-        with open(path, newline="") as csvfile:
-            reader = csv.DictReader(csvfile, delimiter=delimiter, fieldnames=fieldnames)
-            return Clumper(list(it.islice(reader, 0, n)))
+            for row in self:
+                writer.writerow(row)
 
     def _create_new(self, blob):
         """


Test patch:
diff --git a/tests/data/monopoly.csv b/tests/data/monopoly.csv
index 1a5fe78..f5ea185 100644
--- a/tests/data/monopoly.csv
+++ b/tests/data/monopoly.csv
@@ -1,23 +1,23 @@
-name,rent,house_1,house_2,house_3,house_4,hotel,deed_cost,house_cost,color,tile
-Mediterranean Avenue,2,10,30,90,160,250,60,50,purple,1
-Baltic Avenue,4,20,60,180,320,450,60,50,purple,3
-Oriental Avenue,6,30,90,270,400,550,100,50,light_blue,6
-Vermont Avenue,6,30,90,270,400,550,100,50,light_blue,8
-Connecticut Avenue,8,40,100,300,450,600,120,50,light_blue,9
-St. Charles Place,10,50,150,450,625,750,140,100,pink,11
-States Avenue,10,50,150,450,625,750,140,100,pink,13
-Virginia Avenue,12,60,180,500,700,900,160,100,pink,14
-Tennessee Avenue,14,70,200,550,750,950,180,100,orange,16
-St. James Place,14,70,200,550,750,950,180,100,orange,18
-New York Avenue,16,80,220,600,800,1000,200,100,orange,19
-Kentucky Avenue,18,90,250,700,875,1050,220,150,red,21
-Indiana Avenue,18,90,250,700,875,1050,220,150,red,23
-Illinois Avenue,20,100,300,750,925,1100,240,150,red,24
-Atlantic Avenue,22,110,330,800,975,1150,260,150,yellow,26
-Ventnor Avenue,22,110,330,800,975,1150,260,150,yellow,27
-Marvin Gardens,24,120,360,850,1025,1200,280,150,yellow,29
-Pacific Avenue,26,130,390,900,1100,1275,300,200,green,31
-North Carolina Avenue,26,130,390,900,1100,1275,300,200,green,32
-Pennsylvania Avenue,28,150,450,1000,1200,1400,320,200,green,34
-Park Place,35,175,500,1100,1300,1500,350,200,blue,37
+name,rent,house_1,house_2,house_3,house_4,hotel,deed_cost,house_cost,color,tile
+Mediterranean Avenue,2,10,30,90,160,250,60,50,purple,1
+Baltic Avenue,4,20,60,180,320,450,60,50,purple,3
+Oriental Avenue,6,30,90,270,400,550,100,50,light_blue,6
+Vermont Avenue,6,30,90,270,400,550,100,50,light_blue,8
+Connecticut Avenue,8,40,100,300,450,600,120,50,light_blue,9
+St. Charles Place,10,50,150,450,625,750,140,100,pink,11
+States Avenue,10,50,150,450,625,750,140,100,pink,13
+Virginia Avenue,12,60,180,500,700,900,160,100,pink,14
+Tennessee Avenue,14,70,200,550,750,950,180,100,orange,16
+St. James Place,14,70,200,550,750,950,180,100,orange,18
+New York Avenue,16,80,220,600,800,1000,200,100,orange,19
+Kentucky Avenue,18,90,250,700,875,1050,220,150,red,21
+Indiana Avenue,18,90,250,700,875,1050,220,150,red,23
+Illinois Avenue,20,100,300,750,925,1100,240,150,red,24
+Atlantic Avenue,22,110,330,800,975,1150,260,150,yellow,26
+Ventnor Avenue,22,110,330,800,975,1150,260,150,yellow,27
+Marvin Gardens,24,120,360,850,1025,1200,280,150,yellow,29
+Pacific Avenue,26,130,390,900,1100,1275,300,200,green,31
+North Carolina Avenue,26,130,390,900,1100,1275,300,200,green,32
+Pennsylvania Avenue,28,150,450,1000,1200,1400,320,200,green,34
+Park Place,35,175,500,1100,1300,1500,350,200,blue,37
 Boardwalk,50,200,600,1400,1700,2000,400,200,blue,39
diff --git a/tests/data/null.csv b/tests/data/null.csv
new file mode 100644
index 0000000..6071255
--- /dev/null
+++ b/tests/data/null.csv
@@ -0,0 +1,3 @@
+a,b,c
+1,2,
+2,,4
diff --git a/tests/test_read_write/test_read_csv.py b/tests/test_read_write/test_read_csv.py
index c2c5c5a..df0ee33 100644
--- a/tests/test_read_write/test_read_csv.py
+++ b/tests/test_read_write/test_read_csv.py
@@ -1,6 +1,7 @@
 import pytest
 from itertools import product
 from clumper import Clumper
+from string import ascii_uppercase
 
 
 paths = ["tests/data/monopoly.csv", "https://calmcode.io/datasets/monopoly.csv"]
@@ -89,3 +90,24 @@ def test_read_csv_negative_zero():
     """Test that an error is raised if nrows is zero."""
     with pytest.raises(ValueError):
         Clumper.read_csv("tests/data/monopoly.csv", n=0)
+
+
+def test_read_csv_nulls():
+    """Test that null rows are discarded."""
+    assert Clumper.read_csv("tests/data/null.csv").equals(
+        Clumper([{"a": "1", "b": "2"}, {"a": "2", "c": "4"}])
+    )
+
+
+def test_fieldnames_less_than():
+    """Test number of keys returned when an incomplete number of fields is supplied."""
+    reader = Clumper.read_csv("tests/data/monopoly.csv", fieldnames=["A", "B", "C"])
+    assert set(reader.keys()) == {"A", "B", "C", None}
+    assert len(reader.keys()) == len(["A", "B", "C"]) + 1
+
+
+def test_fieldnames_greater_than():
+    """Test value of last key in row is None when an excess number of fields is supplied."""
+    fieldnames = list(ascii_uppercase)[:12]
+    reader = Clumper.read_csv("tests/data/monopoly.csv", fieldnames=fieldnames)
+    assert reader.select(fieldnames[-1]).head(1).equals([{"L": None}])
diff --git a/tests/test_read_write/test_write_csv.py b/tests/test_read_write/test_write_csv.py
new file mode 100644
index 0000000..71935c8
--- /dev/null
+++ b/tests/test_read_write/test_write_csv.py
@@ -0,0 +1,112 @@
+import pytest
+from clumper import Clumper
+
+
+def temp_file(tmp_path):
+    """helper function to create temporary file."""
+    d = tmp_path / "sub"
+    d.mkdir()
+    path = d / "nulls.csv"
+    return path
+
+
+def test_write_csv(tmp_path):
+    path = temp_file(tmp_path)
+    Clumper.read_csv("tests/data/monopoly.csv").write_csv(path)
+    reader = Clumper.read_csv(path)
+    assert Clumper.read_csv("tests/data/monopoly.csv").collect() == reader.collect()
+
+
+def test_write_csv_append_mode(tmp_path):
+    """Test that the append mode works."""
+    path = temp_file(tmp_path)
+    Clumper.read_csv("tests/data/monopoly.csv").write_csv(path)
+    Clumper.read_csv("tests/data/monopoly.csv").write_csv(path, mode="a")
+    reader = Clumper.read_csv(path)
+    # the header is appended as well, hence the +1
+    assert len(reader) == (2 * len(Clumper.read_csv("tests/data/monopoly.csv"))) + 1
+
+
+@pytest.fixture()
+def data():
+    return [
+        {"One": "a", "Two": "1", "Three": "one"},
+        {"One": "b", "Two": "2", "Three": "two"},
+        {"One": "", "Two": "3", "Three": "three"},
+        {"One": "d", "Two": "4", "Three": ""},
+        {"One": "e", "Two": "5", "Three": "five"},
+        {"One": "", "Two": "6", "Three": ""},
+        {"One": "g", "Two": "7", "Three": "seven"},
+    ]
+
+
+def test_write_empty_csv(data, tmp_path):
+    """Test that null cells are exported correctly as empty strings"""
+    path = temp_file(tmp_path)
+    Clumper(data).write_csv(path)
+    reader = Clumper.read_csv(path, na_values="ignore")
+    assert Clumper(data).collect() == reader.collect()
+
+
+def test_write_csv_fieldnames(data, tmp_path):
+    """Test that fieldnames of Clumper match the headers in the exported csv file"""
+    path = temp_file(tmp_path)
+    Clumper(data).write_csv(path)
+    reader = Clumper.read_csv(path)
+    assert not set(Clumper(data).keys()).difference(reader.keys())
+
+
+def test_write_csv_n_positive(data, tmp_path):
+    """Test that the correct number of rows is exported"""
+    path = temp_file(tmp_path)
+    Clumper(data).head(n=10).write_csv(path)
+    reader = Clumper.read_csv(path, na_values="ignore")
+    assert Clumper(data).head(n=10).collect() == reader.collect()
+
+
+def test_write_missing_keys(tmp_path):
+    """Test that function works with missing keys."""
+    missing_key = [{"a": "1", "b": "2"}, {"c": "3"}]
+    path = temp_file(tmp_path)
+    Clumper(missing_key).write_csv(path)
+    reader = Clumper.read_csv(path)
+    assert Clumper(missing_key).equals(reader)
+
+
+data_type = [
+    [{"a": 1, "b": 2}, {"a": 2, "b": 3, "c": 4}, {"a": 1, "b": 6}],
+    [{"a": 1, "b": 2}, {"c": 3}],
+    [{"a": 1, "b": 2}, {"a": 3, "b": 3}, {"a": 2, "b": 1}],
+    [
+        {"Name": "Sam", "Age": 34, "City": "Sydney"},
+        {"Name": "Ade", "Age": 31, "City": "Lagos"},
+        {"Name": "Uche", "Age": 16, "City": "Abuja"},
+        {"Name": "Maleek", "Age": 32, "City": "Kano"},
+        {"Name": "Ragnar", "Age": 33, "City": "Reykjavik"},
+        {"Name": "Zeus", "Age": 35, "City": "Athens"},
+    ],
+    [
+        {"V4": "A", "V1": 0, "V2": 3},
+        {"V4": "A", "V1": 0, "V2": 4},
+        {"V4": "A", "V1": 0, "V2": 7},
+        {"V4": "B", "V1": 4, "V2": 0},
+        {"V4": "B", "V1": 4, "V2": 8},
+        {"V4": "B", "V1": 1, "V2": 5},
+        {"V4": "C", "V1": 4, "V2": 5},
+        {"V4": "C", "V1": 1, "V2": 0},
+        {"V4": "C", "V1": 1, "V2": 9},
+    ],
+]
+
+data_type = tuple(
+    zip(data_type, ["int", "int", "int", {"Age": "int"}, {"V1": "int", "V2": "int"}])
+)
+
+
+@pytest.mark.parametrize("dtype_data,dtype", data_type)
+def test_read_csv(dtype_data, dtype, tmp_path):
+    """Test that the correct dtype is returned when dtype argument is not None."""
+    path = temp_file(tmp_path)
+    Clumper(dtype_data).write_csv(path)
+    reader = Clumper.read_csv(path, dtype=dtype)
+    assert Clumper(dtype_data).equals(reader)
diff --git a/tests/test_read_write/test_write_json.py b/tests/test_read_write/test_write_json.py
new file mode 100644
index 0000000..6dc45dc
--- /dev/null
+++ b/tests/test_read_write/test_write_json.py
@@ -0,0 +1,30 @@
+import pytest
+import os
+from clumper import Clumper
+
+
+def test_local_write_exists(tmp_path):
+    """Test that an error is raised if the written file JSON doesn't exists."""
+    path = str(tmp_path / "pokemon_copy.json")
+    clump = Clumper.read_json("tests/data/pokemon.json")
+    clump.write_json(path)
+    assert os.path.exists(path)
+
+
+def test_local_read_write_content_same(tmp_path):
+    """Test that an error is raised if the written JSON file is not the same as what is read locally"""
+    path = str(tmp_path / "pokemon_copy.json")
+    writer = Clumper.read_json("tests/data/pokemon.json")
+    writer.write_json(path)
+    reader = Clumper.read_json(path)
+    assert reader.collect() == writer.collect()
+
+
+@pytest.mark.parametrize("lines, expected", [(None, 4), (1, 1), (2, 2), (5, 4)])
+def test_local_read_write_same_lines(tmp_path, lines, expected):
+    """Test that an error is raised if the locally written files has the same number of lines as expected"""
+    path = str(tmp_path / "pokemon_copy.json")
+    writer = Clumper.read_json("tests/data/pokemon.json", lines)
+    writer.write_json(path)
+    reader = Clumper.read_json(path)
+    assert len(reader) == len(writer)
diff --git a/tests/test_read_write/test_write_jsonl.py b/tests/test_read_write/test_write_jsonl.py
new file mode 100644
index 0000000..8d83450
--- /dev/null
+++ b/tests/test_read_write/test_write_jsonl.py
@@ -0,0 +1,32 @@
+import pytest
+import os
+from clumper import Clumper
+
+
+def test_local_write_exists(tmp_path):
+    """Test that an error is raised if the written file JSONL doesn't exists."""
+    path = str(tmp_path / "cards_copy.jsonl")
+    clump = Clumper.read_jsonl("tests/data/cards.jsonl")
+    clump.write_jsonl(path)
+    assert os.path.exists(path)
+
+
+def test_local_read_write_content_same(tmp_path):
+    """Test that the written JSONL file is the same as what is read locally"""
+    path = str(tmp_path / "cards_copy.jsonl")
+    writer = Clumper.read_jsonl("tests/data/cards.jsonl")
+    writer.write_jsonl(path)
+    reader = Clumper.read_jsonl(path)
+    assert reader.collect() == writer.collect()
+
+
+@pytest.mark.parametrize("lines, expected", [(1, 1), (2, 2), (5, 4)])
+def test_local_read_write_same_lines(tmp_path, lines, expected):
+    """Test that the locally written files has the same number of lines as expected"""
+    path = tmp_path / "cards_copy.jsonl"
+
+    writer = Clumper.read_jsonl("tests/data/cards.jsonl", lines)
+    writer.write_jsonl(path)
+
+    reader = Clumper.read_jsonl(str(path))
+    assert len(reader) == len(writer)