From 9d23044f288a3d34dae855da9c9eaf5a6a339e5c Mon Sep 17 00:00:00 2001
From: Auke Willem Oosterhoff <1565144+OrangeTux@users.noreply.github.com>
Date: Mon, 17 Jan 2022 08:02:35 +0100
Subject: [PATCH] Fix bug in `remove_nones()` when processing `str` (#290)

`remove_nones()` only works with `list`s and `dict`s, not with other
iterables like `str`. However, the function didn't properly checked its
input and would crash when passed a `list` of `str`s.

Fixes: #289
---
 ocpp/charge_point.py       | 22 +++++++++-------------
 tests/test_charge_point.py | 19 ++++++++++++++++++-
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/ocpp/charge_point.py b/ocpp/charge_point.py
index cbdefb89b..364b1f35c 100644
--- a/ocpp/charge_point.py
+++ b/ocpp/charge_point.py
@@ -4,6 +4,7 @@
 import uuid
 import re
 import time
+from typing import Union, List, Dict
 from dataclasses import asdict
 
 from ocpp.routing import create_route_map
@@ -70,19 +71,14 @@ def snake_to_camel_case(data):
     return data
 
 
-def remove_nones(dict_to_scan):
-    new_dict = {}
-    for k, v in dict_to_scan.items():
-        if isinstance(v, dict):
-            v = remove_nones(v)
-        if isinstance(v, list):
-            new_list = []
-            for item in v:
-                new_list.append(remove_nones(item))
-            v = new_list
-        if v is not None:
-            new_dict[k] = v
-    return new_dict
+def remove_nones(data: Union[List, Dict]) -> Union[List, Dict]:
+    if isinstance(data, dict):
+        return {k: remove_nones(v) for k, v in data.items() if v is not None}
+
+    elif isinstance(data, list):
+        return [remove_nones(v) for v in data if v is not None]
+
+    return data
 
 
 class ChargePoint:
diff --git a/tests/test_charge_point.py b/tests/test_charge_point.py
index 9db1107d2..c966d4175 100644
--- a/tests/test_charge_point.py
+++ b/tests/test_charge_point.py
@@ -3,7 +3,8 @@
 import pytest
 from ocpp.v20 import ChargePoint as cp
 from ocpp.routing import on, create_route_map
-from ocpp.v16.call import BootNotificationPayload, MeterValuesPayload
+from ocpp.v16.call import (BootNotificationPayload, MeterValuesPayload,
+                           GetConfigurationPayload)
 from ocpp.v16.enums import Action
 from ocpp.v16.datatypes import MeterValue, SampledValue
 from ocpp.v201.call import SetNetworkProfilePayload
@@ -148,3 +149,19 @@ def test_nested_list_remove_nones():
 
     payload = asdict(payload)
     assert expected_payload == remove_nones(payload)
+
+
+def test_remove_nones_with_list_of_strings():
+    """ Make sure that `remove_nones()` doesn't crash when it encounters an
+    iterable other than a list or dict. See
+    https://github.com/mobilityhouse/ocpp/issues/289.
+    """
+    payload = asdict(
+        GetConfigurationPayload(
+            key=["ClockAlignedDataInterval", "ConnectionTimeOut"]
+        )
+    )
+
+    assert remove_nones(payload) == {
+        'key': ['ClockAlignedDataInterval', 'ConnectionTimeOut']
+    }