add arugment for primitive type to create_object

shogun-toolbox · Mar 19, 2018 · ae8a7b9 · ae8a7b9
1 parent 278ad1b
commit ae8a7b9
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 6 deletions.
diff --git a/src/shogun/base/class_list.h b/src/shogun/base/class_list.h
@@ -7,9 +7,11 @@
 #ifndef __SG_CLASS_LIST_H__
 #define __SG_CLASS_LIST_H__
 
+#include <shogun/base/SGObject.h>
 #include <shogun/lib/config.h>
 
 #include <shogun/lib/DataType.h>
+#include <shogun/lib/ShogunException.h>
 
 #include <shogun/io/SGIO.h>
 
@@ -37,19 +39,24 @@ namespace shogun {
 	 *
 	 */
 	template <class T>
-	T* create_object(const char* name)
+	T* create_object(const char* name, EPrimitiveType pt=PT_NOT_GENERIC) throw (ShogunException)
 	{
-		auto* object = create(name, PT_NOT_GENERIC);
+		auto* object = create(name, pt);
 		if (!object)
 		{
-			SG_SERROR("No such class %s", name);
+			SG_SERROR("Class %s with primitive type %s does not exist.\n", name, ptype(pt).c_str());
 		}
-		auto* cast = dynamic_cast<T*>(object);
-		if (!cast)
+		T* cast = nullptr;
+		try
+		{
+			cast = object->as<T>();
+		}
+		catch (const ShogunException& e)
 		{
 			delete_object(object);
-			SG_SERROR("Type mismatch");
+			throw e;
 		}
+
 		cast->ref();
 		return cast;
 	}

diff --git a/tests/unit/base/create_unittest.cc b/tests/unit/base/create_unittest.cc
@@ -2,6 +2,7 @@
 #include <shogun/kernel/Kernel.h>
 #include <shogun/machine/KernelMachine.h>
 #include <shogun/machine/Machine.h>
+#include <shogun/features/DenseFeatures.h>
 
 #include <gtest/gtest.h>
 
@@ -17,6 +18,16 @@ TEST(CreateObject,create_wrong_type)
     EXPECT_THROW(create_object<CMachine>("GaussianKernel"), ShogunException);
 }
 
+TEST(CreateObject,create_wrong_ptype)
+{
+    EXPECT_THROW(create_object<CMachine>("GaussianKernel", PT_FLOAT64), ShogunException);
+}
+
+TEST(CreateObject,create_wrong_ptype2)
+{
+    EXPECT_THROW(create_object<CMachine>("DenseFeatures"), ShogunException);
+}
+
 TEST(CreateObject,create_wrong_type_wrong_name)
 {
     EXPECT_THROW(create_object<CMachine>("GoussianKernel"), ShogunException);
@@ -27,6 +38,7 @@ TEST(CreateObject,create)
     auto* obj = create_object<CKernel>("GaussianKernel");
     EXPECT_TRUE(obj != nullptr);
     EXPECT_TRUE(dynamic_cast<CKernel*>(obj) != nullptr);
+    EXPECT_EQ(obj->get_generic(), PT_NOT_GENERIC);
     delete obj;
 }