Add giatools.pandas and tests

giatools/__init__.py

@@ -1,5 +1,5 @@
 VERSION_MAJOR = 0
 VERSION_MINOR = 3
-VERSION_PATCH = 0
+VERSION_PATCH = 1
 
 VERSION = '%d.%d%s' % (VERSION_MAJOR, VERSION_MINOR, '.%d' % VERSION_PATCH if VERSION_PATCH > 0 else '')

giatools/pandas.py

@@ -0,0 +1,33 @@
+from typing import (
+    Any,
+    Iterable,
+    overload,
+)
+
+try:
+    import pandas as pd
+
+    @overload
+    def find_column(df: pd.DataFrame, candidates: Iterable[str]) -> str: ...  # type: ignore
+
+except ImportError:
+    pd = None
+
+    @overload
+    def find_column(df: Any, candidates: Iterable[str]) -> str: ...
+
+
+def find_column(df: Any, candidates: Iterable[str]) -> str:
+    """
+    Returns the column name present in `df` and the list of `candidates`.
+
+    Raises:
+        KeyError: If there is no candidate column name present in `df`, or more than one.
+    """
+    intersection = frozenset(df.columns) & frozenset(candidates)
+    if len(intersection) == 0:
+        raise KeyError(f'No such column: {", ".join(candidates)}')
+    elif len(intersection) > 1:
+        raise KeyError(f'The column names {", ".join(intersection)} are ambiguous')
+    else:
+        return next(iter(intersection))

tests/test_pandas.py

@@ -0,0 +1,31 @@
+import importlib
+import unittest
+
+import pandas as pd
+
+import giatools.pandas
+
+
+class find_column(unittest.TestCase):
+
+    def setUp(self):
+        self.df = pd.DataFrame({'a': [1, 2], 'b': [3, 4]})
+
+    def test(self):
+        self.assertEqual(giatools.pandas.find_column(self.df, ['a', 'c']), 'a')
+
+    def test__missing(self):
+        with self.assertRaises(KeyError):
+            giatools.pandas.find_column(self.df, ['c'])
+
+    def test__ambiguous(self):
+        with self.assertRaises(KeyError):
+            giatools.pandas.find_column(self.df, ['a', 'b'])
+
+
+@unittest.mock.patch.dict('sys.modules', {'pandas': None})
+class no_pandas(unittest.TestCase):
+
+    def test(self):
+        importlib.reload(giatools.pandas)
+        assert giatools.pandas.pd is None