From fc9b22e796f95ee0cc24529e9a3516ce390e83a3 Mon Sep 17 00:00:00 2001
From: Dun Liang <randonlang@gmail.com>
Date: Tue, 28 Apr 2020 21:50:10 +0800
Subject: [PATCH] jt.sum -> nanostring(add)

---
 python/jittor/test/test_nano_string.py | 8 ++++++++
 src/misc/nano_string.cc                | 1 +
 src/pyjt/py_converter.h                | 7 +++++++
 3 files changed, 16 insertions(+)

diff --git a/python/jittor/test/test_nano_string.py b/python/jittor/test/test_nano_string.py
index 2feb43cd..81eb1a22 100644
--- a/python/jittor/test/test_nano_string.py
+++ b/python/jittor/test/test_nano_string.py
@@ -44,6 +44,14 @@ class TestNanoString(unittest.TestCase):
         assert str(jt.NanoString(np.float64)) == "float64"
         assert str(jt.NanoString(np.int8)) == "int8"
         assert str(jt.NanoString(np.array([1,2,3]).dtype)) == "int64"
+
+        assert str(jt.NanoString(jt.float)) == "float"
+        assert str(jt.NanoString(jt.float32)) == "float32"
+        assert str(jt.NanoString(jt.float64)) == "float64"
+        assert str(jt.NanoString(jt.int8)) == "int8"
+        assert str(jt.NanoString(jt.array([1,2,3]).dtype)) == "int64"
+        assert str(jt.NanoString(jt.sum)) == "add"
+
         def get_error_str(call):
             es = ""
             try:
diff --git a/src/misc/nano_string.cc b/src/misc/nano_string.cc
index 3585c78a..a64ae48a 100644
--- a/src/misc/nano_string.cc
+++ b/src/misc/nano_string.cc
@@ -147,6 +147,7 @@ static void init_ns() {
     #define INIT_NS(T) func(#T, ns_##T);
     FOR_ALL_NS(INIT_NS);
     ASSERT(NanoString::__ns_to_string.size()<=(1<<NanoString::_index_nbits));
+    NanoString::__string_to_ns["sum"] = ns_add;
     LOGvv << "init __string_to_ns" << NanoString::__string_to_ns;
     LOGvv << "init __ns_to_string" << NanoString::__ns_to_string;
 }
diff --git a/src/pyjt/py_converter.h b/src/pyjt/py_converter.h
index 824f1ca6..dde2897e 100644
--- a/src/pyjt/py_converter.h
+++ b/src/pyjt/py_converter.h
@@ -161,6 +161,8 @@ DEF_IS(NanoString, bool) is_type(PyObject* obj) {
     return Py_TYPE(obj) == &PyjtNanoString ||
         PyUnicode_CheckExact(obj) ||
         PyType_CheckExact(obj) ||
+        // jt.float.__name__
+        PyCallable_Check(obj) ||
         // numpy.dtype.type
         PyObject_HasAttrString(obj, "type");
 }
@@ -180,6 +182,11 @@ DEF_IS(NanoString, T) from_py_object(PyObject* obj) {
     // PyType
     if (PyType_CheckExact(obj))
         return T(_PyType_Name((PyTypeObject *)obj));
+    // jt.float.__name__
+    if (PyCallable_Check(obj)) {
+        PyObjHolder t(PyObject_GetAttrString(obj, "__name__"));
+        return T(PyUnicode_AsUTF8(t.obj));
+    }
     PyObjHolder t(PyObject_GetAttrString(obj, "type"));
     CHECK(PyType_CheckExact(t.obj)) << "Not a valid type:" << t.obj;
     return T(_PyType_Name((PyTypeObject *)t.obj));