Refactor Python string utility function

torch/csrc/Module.cpp

@@ -7,6 +7,8 @@
 #include <libshm.h>
 #include <TH/TH.h>
 
+#include "torch/csrc/utils/python_strings.h"
+
 #ifdef WITH_CUDNN
 #include "cudnn/Module.h"
 #endif
@@ -73,13 +75,9 @@ static PyObject * THPModule_initNames(PyObject *self, PyObject *arg)
 
     THPObjectPtr module_name = PyObject_GetAttrString(obj, "__module__");
     if (!module_name) return NULL;
-#if PY_MAJOR_VERSION == 2
-    THPUtils_assert(PyString_Check(module_name.get()), "expected __module__ to be a string");
-    std::string name = PyString_AS_STRING(module_name.get());
-#else
-    THPUtils_assert(PyUnicode_Check(module_name.get()), "expected __module__ to be a string");
-    std::string name = PyUnicode_AsUTF8(module_name.get());
-#endif
+    THPUtils_assert(THPUtils_checkString(module_name.get()),
+        "expected __module__ to be a string");
+    std::string name = THPUtils_unpackString(module_name.get());
     names.push_back(name + "." + type->tp_name);
     type->tp_name = names.back().c_str();
   }
@@ -89,15 +87,13 @@ static PyObject * THPModule_initNames(PyObject *self, PyObject *arg)
 static bool THPModule_assignStateless(PyObject *self)
 {
 #define INIT_STATELESS(type)                                                   \
-  stateless = PyObject_Call((PyObject*)&TH_CONCAT_2(type, TensorStatelessType), arg, NULL); \
+  stateless = PyObject_CallFunctionObjArgs((PyObject*)&TH_CONCAT_2(type, TensorStatelessType), NULL); \
   if (!stateless) {                                                            \
-    THPUtils_setError("stateless method initialization error");                \
     return false;                                                              \
   }                                                                            \
   if (PyObject_SetAttrString(TH_CONCAT_3(THP,type,TensorClass), THP_STATELESS_ATTRIBUTE_NAME, stateless) == -1) { \
-    THPUtils_setError("stateless method initialization error (on assignment)");\
+    return false;                                                              \
   }
-  PyObject *arg = PyTuple_New(0);
   PyObject *stateless;
   INIT_STATELESS(Double);
   INIT_STATELESS(Float);
@@ -107,23 +103,25 @@ static bool THPModule_assignStateless(PyObject *self)
   INIT_STATELESS(Short);
   INIT_STATELESS(Char);
   INIT_STATELESS(Byte);
-  Py_DECREF(arg);
   return true;
 #undef INIT_STATELESS
 }
 //
 // Callback for python part. Used for additional initialization of python classes
 static PyObject * THPModule_initExtension(PyObject *self, PyObject *shm_manager_path)
 {
-  if (!THPUtils_checkBytes(shm_manager_path)) {
+  HANDLE_TH_ERRORS
+  if (!THPUtils_checkString(shm_manager_path)) {
     THPUtils_setError("initialization error - expected bytes/string object as shm_manager_path!");
     return NULL;
   }
-  libshm_init(THPUtils_bytesAsString(shm_manager_path));
+  std::string path = THPUtils_unpackString(shm_manager_path);
+  libshm_init(path.c_str());
   if (!THPModule_loadClasses(self))         return NULL;
   if (!THPModule_assignStateless(self))     return NULL;
   if (!THPAutograd_initFunctions(self))     return NULL;
-  return PyBool_FromLong(true);
+  Py_RETURN_NONE;
+  END_HANDLE_TH_ERRORS
 }
 
 static PyObject * THPModule_getNumThreads(PyObject *module)
@@ -429,22 +427,6 @@ PyObject *THPModule_safeCall(PyObject *_unused, PyObject *args, PyObject *kwargs
   return result;
 }
 
-static std::string parseString(PyObject *obj)
-{
-  if (PyBytes_Check(obj)) {
-    return std::string(PyBytes_AS_STRING(obj));
-#if PY_MAJOR_VERSION == 3
-  } else if (PyUnicode_Check(obj)) {
-    return std::string(PyUnicode_AsUTF8(obj));
-#else
-  } else if (PyUnicode_Check(obj)) {
-    THPObjectPtr utf8 = PyUnicode_AsUTF8String(obj);
-    return std::string(PyBytes_AS_STRING(utf8.get()));
-#endif
-  }
-  return "<invalid string>";
-}
-
 PyObject *THPModule_addDocStr(PyObject *_unused, PyObject *args)
 {
   // adds a __doc__ string to a function, similar to numpy's arr_add_docstring
@@ -455,8 +437,11 @@ PyObject *THPModule_addDocStr(PyObject *_unused, PyObject *args)
     return NULL;
   }
 
-  all_docs.push_back(parseString(doc_obj));
-  const char* doc_str = all_docs.back().c_str();
+  const char* doc_str = "<invalid string>";
+  if (THPUtils_checkString(doc_obj)) {
+    all_docs.push_back(THPUtils_unpackString(doc_obj));
+    doc_str = all_docs.back().c_str();
+  }
 
   if (Py_TYPE(obj) == &PyCFunction_Type) {
     PyCFunctionObject* f = (PyCFunctionObject *)obj;
@@ -499,7 +484,6 @@ extern PyObject * THCPModule_seedAll(PyObject *_unused);
 extern PyObject * THCPModule_initialSeed(PyObject *_unused);
 extern PyObject * THCPModule_cudaHostAllocator(PyObject *_unused);
 extern PyObject * THCPModule_cudaSynchronize(PyObject *_unused);
-extern PyObject * THCPModule_getLibPath(PyObject *_unused);
 extern PyObject * THCPModule_cudaSleep(PyObject *_unused, PyObject *cycles);
 extern PyObject * THCPModule_cudaLockMutex(PyObject *module);
 extern PyObject * THCPModule_cudaUnlockMutex(PyObject *module);
@@ -532,7 +516,6 @@ static PyMethodDef TorchMethods[] = {
   {"_cuda_initialSeed", (PyCFunction)THCPModule_initialSeed,      METH_NOARGS,  NULL},
   {"_cuda_cudaHostAllocator", (PyCFunction)THCPModule_cudaHostAllocator, METH_NOARGS, NULL},
   {"_cuda_synchronize", (PyCFunction)THCPModule_cudaSynchronize, METH_NOARGS, NULL},
-  {"_cuda_getLibPath", (PyCFunction)THCPModule_getLibPath, METH_NOARGS, NULL},
   {"_cuda_sleep", (PyCFunction)THCPModule_cudaSleep, METH_O, NULL},
   {"_cuda_sparse_init",  (PyCFunction)THCSPModule_initExtension,    METH_NOARGS,  NULL},
   {"_cuda_lock_mutex",   (PyCFunction)THCPModule_cudaLockMutex,   METH_NOARGS,  NULL},

torch/csrc/Size.cpp

@@ -1,6 +1,7 @@
 #include "Size.h"
 
 #include <string>
+#include "torch/csrc/utils/python_strings.h"
 #include "THP.h"
 
 PyObject* THPSizeClass = NULL;
@@ -47,11 +48,7 @@ static PyObject * THPSize_repr(THPSize *self)
     repr += std::to_string(PyLong_AsLong(PyTuple_GET_ITEM(self, i)));
   }
   repr += "])";
-#if PY_MAJOR_VERSION == 2
-  return PyString_FromString(repr.c_str());
-#else
-  return PyUnicode_FromString(repr.c_str());
-#endif
+  return THPUtils_packString(repr);
 }
 
 extern PyTypeObject THPSizeType;

torch/csrc/autograd/python_hook.cpp

@@ -6,6 +6,7 @@
 #include "torch/csrc/autograd/python_variable.h"
 #include "torch/csrc/utils/auto_gil.h"
 #include "torch/csrc/utils/object_ptr.h"
+#include "torch/csrc/utils/python_strings.h"
 #include "torch/csrc/Exceptions.h"
 #include <THPP/THPP.h>
 
@@ -184,15 +185,8 @@ static void check_single_result(PyObject* _original, PyObject* _result, PyObject
 
 static std::string hook_name(PyObject* hook) {
   THPObjectPtr name = PyObject_GetAttrString(hook, "__name__");
-#if PY_MAJOR_VERSION == 2
-  if (name && PyString_Check(name.get())) {
-    return std::string(PyString_AS_STRING(name.get()));
+  if (name && THPUtils_checkString(name.get())) {
+    return THPUtils_unpackString(name.get());
   }
-#else
-  if (name && PyUnicode_Check(name.get())) {
-    THPObjectPtr tmp = PyUnicode_AsASCIIString(name.get());
-    return std::string(PyBytes_AS_STRING(tmp.get()));
-  }
-#endif
   return "<unknown>";
 }

torch/csrc/cuda/Module.cpp

@@ -259,19 +259,6 @@ PyObject * THCPModule_cudaUnlockMutex(PyObject *module)
   Py_RETURN_NONE;
 }
 
-PyObject * THCPModule_getLibPath(PyObject *_unused)
-{
-#define _STR(x) #x
-#define STR(x) _STR(x)
-#if PY_MAJOR_VERSION == 2
-  return PyString_FromString(STR(CUDA_LIB_PATH));
-#else
-  return PyUnicode_FromString(STR(CUDA_LIB_PATH));
-#endif
-#undef STR
-#undef _STR
-}
-
 ////////////////////////////////////////////////////////////////////////////////
 // Cuda module initialization
 ////////////////////////////////////////////////////////////////////////////////

torch/csrc/distributed/Module.cpp

@@ -42,43 +42,27 @@ static bool THDPModule_loadClasses(PyObject *module_dict)
 static std::unordered_map<PyObject*, THDReduceOp> obj2reduceop;
 static std::unordered_map<PyObject*, THDGroup> obj2group;
 
-static THPObjectPtr _ensureBytes(PyObject *obj)
-{
-#if PY_MAJOR_VERSION == 2
-  if (PyString_Check(obj)) {
-#elif PY_MAJOR_VERSION == 3
-  if (PyBytes_Check(obj)) {
-#endif
-    Py_INCREF(obj);
-    return obj;
-  }
-  if (PyUnicode_Check(obj)) {
-    return PyUnicode_AsASCIIString(obj);
-  }
-  return NULL;
-}
-
-PyObject* THDPModule_initProcessGroup(PyObject *_unused, PyObject *_backend)
+PyObject* THDPModule_initProcessGroup(PyObject *_unused, PyObject *backend)
 {
   HANDLE_TH_ERRORS
-  THPObjectPtr backend_bytes = _ensureBytes(_backend);
-  THPUtils_assert(backend_bytes, "backend argument has to be a string/bytes "
-      "object, but got %s", THPUtils_typename(_backend));
-  char *backend_name = THPUtils_bytesAsString(backend_bytes.get());
+  THPUtils_assert(THPUtils_checkString(backend),
+      "backend argument has to be a string/bytes object, but got %s",
+      THPUtils_typename(backend));
+  std::string backend_name = THPUtils_unpackString(backend);
   THDChannelType channel_type = name2channel_type.at(backend_name);
   THPUtils_assert(THDProcessGroupInit(channel_type), "failed to initialize "
       "distributed library (THD)");
   Py_RETURN_NONE;
   END_HANDLE_TH_ERRORS
 }
 
-PyObject* THDPModule_initMasterWorker(PyObject *_unused, PyObject *_backend)
+PyObject* THDPModule_initMasterWorker(PyObject *_unused, PyObject *backend)
 {
   HANDLE_TH_ERRORS
-  THPObjectPtr backend_bytes = _ensureBytes(_backend);
-  THPUtils_assert(backend_bytes, "backend argument has to be a string/bytes "
-      "object, but got %s", THPUtils_typename(_backend));
-  char *backend_name = THPUtils_bytesAsString(backend_bytes.get());
+  THPUtils_assert(THPUtils_checkString(backend),
+      "backend argument has to be a string/bytes object, but got %s",
+      THPUtils_typename(backend));
+  std::string backend_name = THPUtils_unpackString(backend);
   THDChannelType channel_type = name2channel_type.at(backend_name);
   THPUtils_assert(THDMasterWorkerInit(channel_type), "failed to initialize "
       "distributed library (THD)");

torch/csrc/generic/StorageSharing.cpp

@@ -104,10 +104,10 @@ static PyObject * THPStorage_(shareFilename)(THPStorage *self)
     ctx = (libshm_context*)storage->allocatorContext;
   }
 
-  THPObjectPtr manager_handle = THPUtils_bytesFromString(ctx->manager_handle);
+  THPObjectPtr manager_handle = PyBytes_FromString(ctx->manager_handle);
   if (!manager_handle) return NULL;
   THPObjectPtr storage_handle =
-    THPUtils_bytesFromString(THMapAllocatorContext_filename(ctx->th_context));
+    PyBytes_FromString(THMapAllocatorContext_filename(ctx->th_context));
   if (!storage_handle) return NULL;
   THPObjectPtr size = PyLong_FromLong(storage->size);
   if (!size) return NULL;
@@ -124,20 +124,21 @@ static PyObject * THPStorage_(shareFilename)(THPStorage *self)
 static PyObject * THPStorage_(newSharedFilename)(PyObject *_unused, PyObject *args)
 {
   HANDLE_TH_ERRORS
+  THPUtils_assert(PyTuple_GET_SIZE(args) == 3, "tuple of 3 items expected");
   PyObject *_manager_handle = PyTuple_GET_ITEM(args, 0);
   PyObject *_object_handle = PyTuple_GET_ITEM(args, 1);
   PyObject *_size = PyTuple_GET_ITEM(args, 2);
-  if (!THPUtils_checkBytes(_manager_handle) || !THPUtils_checkBytes(_object_handle) || !THPUtils_checkLong(_size)) {
+  if (!PyBytes_Check(_manager_handle) || !PyBytes_Check(_object_handle) || !THPUtils_checkLong(_size)) {
     THPUtils_invalidArguments(args, NULL, "_new_shared in file system mode", 1,
         "a handle (string/bytes) and storage size (int)");
     return NULL;
   }
-  const char *manager_handle = THPUtils_bytesAsString(_manager_handle);
-  const char *object_handle = THPUtils_bytesAsString(_object_handle);
+  const char *manager_handle = PyBytes_AS_STRING(_manager_handle);
+  const char *object_handle = PyBytes_AS_STRING(_object_handle);
   long size = THPUtils_unpackLong(_size);
-
-  libshm_context *ctx = libshm_context_new(manager_handle, object_handle,
-          TH_ALLOCATOR_MAPPED_SHAREDMEM | TH_ALLOCATOR_MAPPED_NOCREATE);
+  int flags = TH_ALLOCATOR_MAPPED_SHAREDMEM |
+              TH_ALLOCATOR_MAPPED_NOCREATE;
+  libshm_context *ctx = libshm_context_new(manager_handle, object_handle, flags);
   return THPStorage_(New)(THStorage_(newWithAllocator)(size,
       &THManagedSharedAllocator, (void*)ctx));
   END_HANDLE_TH_ERRORS
@@ -199,6 +200,7 @@ static PyObject * THPStorage_(shareFd)(THPStorage *self)
 static PyObject * THPStorage_(newSharedFd)(PyObject *_unused, PyObject *args)
 {
   HANDLE_TH_ERRORS
+  THPUtils_assert(PyTuple_GET_SIZE(args) == 2, "tuple of 2 items expected");
   PyObject *_tmp_fd = PyTuple_GET_ITEM(args, 0);
   PyObject *_size = PyTuple_GET_ITEM(args, 1);
   if (!THPUtils_checkLong(_tmp_fd) || !THPUtils_checkLong(_size)) {

torch/csrc/utils.cpp

@@ -6,6 +6,7 @@
 #include <algorithm>
 #include <unordered_map>
 #include "THP.h"
+#include "torch/csrc/utils/python_strings.h"
 
 #include "generic/utils.cpp"
 #include <TH/THGenerateAllTypes.h>
@@ -457,15 +458,6 @@ std::vector<std::string> _tryMatchKwargs(const Option& option,
   return unmatched;
 }
 
-std::string _parseDictKey(PyObject *key_str) {
-#if PY_MAJOR_VERSION == 3
-  THPObjectPtr ascii = PyUnicode_AsASCIIString(key_str);
-  return std::string(PyBytes_AS_STRING(ascii.get()));
-#else
-  return std::string(PyString_AS_STRING(key_str));
-#endif
-}
-
 void THPUtils_invalidArguments(PyObject *given_args, PyObject *given_kwargs,
         const char *function_name, size_t num_options, ...) {
   std::vector<std::string> option_strings;
@@ -493,7 +485,7 @@ void THPUtils_invalidArguments(PyObject *given_args, PyObject *given_kwargs,
     Py_ssize_t pos = 0;
 
     while (PyDict_Next(given_kwargs, &pos, &key, &value)) {
-      kwargs.emplace(_parseDictKey(key), value);
+      kwargs.emplace(THPUtils_unpackString(key), value);
     }
   }
 

torch/csrc/utils.h

@@ -25,17 +25,6 @@
     (throw std::runtime_error("Could not unpack long"), 0))
 #endif
 
-
-#if PY_MAJOR_VERSION == 2
-#define THPUtils_bytesFromString(c_string)   PyString_FromString(c_string)
-#define THPUtils_checkBytes(obj)             PyString_Check(obj)
-#define THPUtils_bytesAsString(obj)          PyString_AS_STRING(obj)
-#else
-#define THPUtils_bytesFromString(c_string)   PyBytes_FromString(c_string)
-#define THPUtils_checkBytes(obj)             PyBytes_Check(obj)
-#define THPUtils_bytesAsString(obj)          PyBytes_AS_STRING(obj)
-#endif
-
 #if PY_MAJOR_VERSION == 2
 #define THPUtils_checkReal_FLOAT(object)                                       \
     (PyFloat_Check(object) || PyLong_Check(object) || PyInt_Check(object))