Added custom dataframe accessors. (#44)

* Added an accessor file to define customer dataframe accessors.

* docstring change.

* Added a test

* Marked another test as flaky

* Added another test to the decode accessor

Author: Kristinn

jupyter/docker/docker_build/00-import.py

@@ -16,6 +16,9 @@
 from picatrix import notebook_init
 import ds4n6_lib as ds
 
+# Add in the accessors to pandas.
+from laceworkjupyter import accessors
+
 # Enable the Picatrix helpers.
 notebook_init.init()
 

jupyter/laceworkjupyter/accessors.py

@@ -0,0 +1,72 @@
+"""
+File that includes custom accessors to pandas.
+
+This is a way to define custom accessors to either
+data frames or series objects.
+"""
+
+import base64
+import urllib.parse
+
+import pandas as pd
+
+
+@pd.api.extensions.register_series_accessor('decode')
+class DecodeAccessor:
+    """
+    Accessor class for decoding data.
+    """
+    def __init__(self, data):
+        self.data = data
+
+    def _decode_string_base64(self, string_value):
+        """
+        Returns decoded base64 string.
+
+        :param str string_value: The base64 decoded string.
+        :return: A decoded string.
+        """
+        decoded_string = base64.b64decode(string_value)
+        try:
+            return decoded_string.decode('utf8')
+        except UnicodeDecodeError:
+            return decoded_string
+
+    def base64(self):
+        """
+        Takes a series with base64 encoded data and decodes it.
+        """
+        return self.data.apply(self._decode_string_base64)
+
+    def base64_altchars(self, altchars):
+        """
+        Takes a series with base64 encoded data and decodes it using altchars.
+
+        :param bytes altchars: A byte-like object of length 2 which specifies
+            the alternative alphabet used instead of the '+' and '/'
+            characters.
+        :return: Decoded Base64 string.
+        """
+        return self.data.apply(
+            lambda x: base64.b64decode(x, altchars=altchars, validate=True))
+
+    def url_unquote(self):
+        """
+        Takes a series with URL encoded characters and unquotes them.
+        """
+        return self.data.apply(urllib.parse.unquote_plus)
+
+
+@pd.api.extensions.register_series_accessor('encode')
+class EncodeAccessor:
+    """
+    Accessor class to encode data.
+    """
+    def __init__(self, data):
+        self.data = data
+
+    def base64(self):
+        """
+        Returns base64 encoded data from a string series.
+        """
+        return self.data.astype(bytes).apply(base64.b64encode)

jupyter/tests/laceworkjupyter/test_accessors.py

@@ -0,0 +1,34 @@
+"""
+Test file for the local accessors.
+"""
+import pandas as pd
+
+from laceworkjupyter import accessors  # noqa: F401
+
+
+def test_decode_accessor():
+    """
+    Tests the decode accessor.
+    """
+    lines = [
+        {'value': 12, 'some_string': 'VGhpcyBpcyBhIHN0cmluZw==', 'uri': 'http://mbl.is/%3Fstuff=r+1%20af'},
+        {'value': 114, 'some_string': 'VGhpcyBpcyBhIGEgc2VjcmV0', 'uri': 'http://mbl.is/%3Fsfi=r+1%20af'},
+    ]
+    frame = pd.DataFrame(lines)
+
+    decoded_series = frame.some_string.decode.base64()
+    discovered_set = set(list(decoded_series.values))
+
+    expected_set = set([
+        'This is a a secret', 'This is a string'])
+
+    assert expected_set == discovered_set
+
+    unquoted_series = frame.uri.decode.url_unquote()
+    unquoted_set = set(list(unquoted_series.values))
+
+    expected_set = set([
+        'http://mbl.is/?stuff=r 1 af',
+        'http://mbl.is/?sfi=r 1 af'])
+
+    assert expected_set == unquoted_set

tests/api/v2/test_agent_access_tokens.py

@@ -30,6 +30,7 @@ def test_agent_access_tokens_api_get(api):
     assert "data" in response.keys()
 
 
+@pytest.mark.flaky_test
 def test_agent_access_tokens_api_get_by_id(api):
     response = api.agent_access_tokens.get()
 
@@ -47,6 +48,7 @@ def test_agent_access_tokens_api_get_by_id(api):
         assert "data" in response.keys()
 
 
+@pytest.mark.flaky_test
 def test_agent_access_tokens_api_search(api):
     assert AGENT_ACCESS_TOKEN_ID is not None
     if AGENT_ACCESS_TOKEN_ID: