diff --git a/src/clj_ml/data.clj b/src/clj_ml/data.clj
index 628f86d..69cf7c5 100644
--- a/src/clj_ml/data.clj
+++ b/src/clj_ml/data.clj
@@ -184,8 +184,8 @@
 
 (defn dataset-class-labels
   "Returns the possible labels for the class attribute"
-  [dataset]
-  (let [class-attr (.classAttribute dataset)
+  [^Instances dataset]
+  (let [^Attribute class-attr (.classAttribute dataset)
         values (.enumerateValues class-attr)]
     (loop [continue (.hasMoreElements values)
            acum {}]
diff --git a/test/clj_ml/data_test.clj b/test/clj_ml/data_test.clj
index ccf3fe8..93db74b 100644
--- a/test/clj_ml/data_test.clj
+++ b/test/clj_ml/data_test.clj
@@ -146,3 +146,9 @@
   (let [ds (make-dataset "test" [:a {:b [:foo :bar]}] [[1 :foo] [2 :bar]])
         _ (dataset-replace-attribute! ds :b (nominal-attribute :b [:baz :shaz]))]
     (is (= [:a {:b [:shaz :baz]}] (dataset-format ds)))))
+
+(deftest dataset-label-helpers
+  (let [ds (make-dataset "test" [:a :b {:c [:d :e]}]
+                         [{:a 1 :b 2 :c :d} [4 5 :e]])]
+    (dataset-set-class ds :c)
+    (is (= {:d 0 :e 1} (dataset-class-labels ds) (dataset-labels-at ds 2)))))