Merge pull request incanter#238 from Gerrrr/feature/sel_string

get-column-id returns keyword version of column-key if convenient
nuroko · Apr 9, 2014 · ffcde29 · ffcde29
2 parents 8de1ccf + d4eb04b
commit ffcde29
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 13 deletions.
diff --git a/modules/incanter-core/src/incanter/core.clj b/modules/incanter-core/src/incanter/core.clj
@@ -1211,19 +1211,27 @@
 
 
 (defn- get-column-id [dataset column-key]
-  (let [headers (:column-names dataset)
-        col-key (if (and
-                     (keyword? column-key) ;; if the given column name is a keyword, and
-                     (not (some #{column-key} headers))) ; a keyword column name wasn't used in the dataset
-                  (name column-key) ;; convert the keyword to a string
-                  column-key) ;; otherwise use the given column key
-        id (if (number? col-key)
-             (if (some #(= col-key %) headers)
-               col-key
-               (nth headers col-key))
-             col-key)]
-    id))
-
+  (let [headers (:column-names dataset)]
+    (cond
+     (and (keyword? column-key) ;; if the given column name is a keyword, and
+          ;; a keyword column name wasn't used in the dataset
+          (not (some #{column-key} headers)))
+     (name column-key) ;; convert the keyword to a string
+
+     (and (string? column-key) ;; if the given column is a string, and
+          ;; this column was't used in the dataset, and
+          (not (some #{column-key} headers))
+          ;; a keyword column name was used in the dataset
+          (some #{(keyword column-key)} headers))
+     ;; convert string to keyword
+     (keyword column-key)
+
+     (and (number? column-key) ;; if the given column name is a number
+          ;; and this number is not in headers
+          (not (some #(= column-key %) headers)))
+     (nth headers column-key) ;; get nth column from headers
+
+     :else column-key)))
 
 (defn- map-get
   ([m k]

diff --git a/modules/incanter-core/test/incanter/core_tests.clj b/modules/incanter-core/test/incanter/core_tests.clj
@@ -38,10 +38,13 @@
 
 (deftest dataset-tests
   (is (= (sel dataset1 :cols :a) [1 4]))
+  (is (= (sel dataset1 :cols "a") [1 4]))
   (is (= (sel dataset1 :all 1) [2 5]))
   (is (= (sel dataset1 :all :b) [2 5]))
   (is (= (sel dataset1 :all [:a :c]) (dataset [:a :c] [[1 3] [4 6]])))
+  (is (= (sel dataset1 :all ["a" "c"]) (dataset [:a :c] [[1 3] [4 6]])))
   (is (= (sel dataset1 :all [:a]) (dataset [:a] [[1] [4]])))
+  (is (= (sel dataset1 :all ["a"]) (dataset [:a] [[1] [4]])))
   (is (= (sel dataset1 :all :all) dataset1))
   (is (= (sel dataset2 :cols :b) [2 5]))
   (is (= (sel dataset2 :cols "c") [3 6]))