From eeaa3ae07cd889800481f89bb90516e88a6f3a9d Mon Sep 17 00:00:00 2001
From: rongxiangsu <rongxiangsu@gmail.com>
Date: Mon, 21 Aug 2023 08:23:15 +0800
Subject: [PATCH] attach_attributes function allows input a list of string

---
 examples/example_turkey_vultures.ipynb | 18 +++++++++---------
 ortega/output.py                       | 12 ++++++++++--
 setup.py                               |  2 +-
 3 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/examples/example_turkey_vultures.ipynb b/examples/example_turkey_vultures.ipynb
index 13368f8..ec9f105 100644
--- a/examples/example_turkey_vultures.ipynb
+++ b/examples/example_turkey_vultures.ipynb
@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 25,
    "id": "f41300e2",
    "metadata": {
     "is_executing": true
@@ -26,7 +26,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 26,
    "id": "ae813970",
    "metadata": {},
    "outputs": [
@@ -48,7 +48,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 27,
    "id": "4e4734c6",
    "metadata": {},
    "outputs": [
@@ -178,7 +178,7 @@
        "4       1.0  11.101806  "
       ]
      },
-     "execution_count": 3,
+     "execution_count": 27,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -208,7 +208,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 28,
    "id": "43ec53ce",
    "metadata": {
     "scrolled": true
@@ -218,8 +218,8 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "CPU times: user 195 ms, sys: 39.5 ms, total: 234 ms\n",
-      "Wall time: 496 ms\n"
+      "CPU times: user 224 ms, sys: 17.6 ms, total: 242 ms\n",
+      "Wall time: 249 ms\n"
      ]
     }
    ],
@@ -243,7 +243,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 29,
    "id": "d651b2ed",
    "metadata": {},
    "outputs": [
@@ -373,7 +373,7 @@
        "4       1.0  11.101806  "
       ]
      },
-     "execution_count": 5,
+     "execution_count": 29,
      "metadata": {},
      "output_type": "execute_result"
     }
diff --git a/ortega/output.py b/ortega/output.py
index f73b1cf..b65f011 100644
--- a/ortega/output.py
+++ b/ortega/output.py
@@ -53,5 +53,13 @@ def compute_interaction_duration(self):
         self.df_interaction_events['duration'] = self.df_interaction_events['duration'].dt.total_seconds().div(60)
         print(datetime.now(), f'Computing interaction duration complete!')
 
-    def attach_attributes(self, col: str, method: str):
-        self.df_all_intersection_pairs = extract_attributes(self.df_all_intersection_pairs, col, method)
+    def attach_attributes(self, col, method: str):
+        if method not in ['mean', 'difference']:
+            raise ValueError("Parameter 'method' must be one of these: ['mean','difference']!")
+        if isinstance(col, str):
+            self.df_all_intersection_pairs = extract_attributes(self.df_all_intersection_pairs, col, method)
+        elif isinstance(col, list):
+            for c in col:
+                self.df_all_intersection_pairs = extract_attributes(self.df_all_intersection_pairs, c, method)
+        else:
+            raise TypeError("Parameter 'col' must be either a string or a list of strings!")
diff --git a/setup.py b/setup.py
index 7666a02..dd53290 100644
--- a/setup.py
+++ b/setup.py
@@ -3,7 +3,7 @@
 if __name__ == "__main__":
     setup(
         name='ortega',
-        version='1.0.12',
+        version='1.0.13',
         author='MOVE lab@UCSB',
         author_email="rongxiangsu@ucsb.edu",
         packages=["ortega"],