Support range and step range for index access

mrkn · mrkn · commit 82742f261e8d · 2017-09-07T11:01:53.000+09:00
diff --git a/ext/pycall/libpython.c b/ext/pycall/libpython.c
@@ -150,6 +150,8 @@ pycall_init_libpython_api_table(VALUE libpython_handle)
   INIT_API_TABLE_ENTRY(PyTuple_GetItem, required);
   INIT_API_TABLE_ENTRY(PyTuple_SetItem, required);
 
+  INIT_API_TABLE_ENTRY(PySlice_New, required);
+
   INIT_API_TABLE_ENTRY(PyIter_Next, required);
 
   INIT_API_TABLE_ENTRY(PyErr_Occurred, required);
diff --git a/ext/pycall/pycall.c b/ext/pycall/pycall.c
@@ -1007,6 +1007,35 @@ pycall_libpython_helpers_m_define_wrapper_method(VALUE mod, VALUE wrapper, VALUE
   return Qnil;
 }
 
+static PyObject *
+pycall_convert_index(VALUE index)
+{
+  PyObject *pyobj;
+
+  if (RB_TYPE_P(index, T_ARRAY) && RARRAY_LEN(index) == 1) {
+    index = RARRAY_AREF(index, 0);
+  }
+  if (RB_TYPE_P(index, T_ARRAY)) {
+    long i, n = RARRAY_LEN(index);
+    pyobj = Py_API(PyTuple_New)(n);
+    for (i = 0; i < n; ++i) {
+      PyObject *pytem = pycall_convert_index(RARRAY_AREF(index, i)); /* New reference */
+      Py_API(PyTuple_SetItem)(pyobj, i, pytem); /* Steal reference */
+    }
+  }
+  else if (rb_obj_is_kind_of(index, rb_cRange)) {
+    pyobj = pycall_pyslice_from_ruby(index); /* New refrence */
+  }
+  else if (pycall_obj_is_step_range(index)) {
+    pyobj = pycall_pyslice_from_ruby(index); /* New refrence */
+  }
+  else {
+    pyobj = pycall_pyobject_from_ruby(index); /* New reference */
+  }
+
+  return pyobj;
+}
+
 static VALUE
 pycall_libpython_helpers_m_getitem(VALUE mod, VALUE pyptr, VALUE key)
 {
@@ -1019,7 +1048,7 @@ pycall_libpython_helpers_m_getitem(VALUE mod, VALUE pyptr, VALUE key)
 
   pyobj = get_pyobj_ptr(pyptr);
 
-  pyobj_key = pycall_pyobject_from_ruby(key);
+  pyobj_key = pycall_convert_index(key);
 
   pyobj_v = Py_API(PyObject_GetItem)(pyobj, pyobj_key);
   if (!pyobj_v) {
@@ -1038,7 +1067,7 @@ pycall_libpython_helpers_m_setitem(VALUE mod, VALUE pyptr, VALUE key, VALUE v)
   int res;
 
   pyobj = check_get_pyobj_ptr(pyptr, NULL);
-  pyobj_key = pycall_pyobject_from_ruby(key);
+  pyobj_key = pycall_convert_index(key);
   pyobj_value = pycall_pyobject_from_ruby(v);
 
   res = Py_API(PyObject_SetItem)(pyobj, pyobj_key, pyobj_value);
@@ -1058,7 +1087,7 @@ pycall_libpython_helpers_m_delitem(VALUE mod, VALUE pyptr, VALUE key, VALUE v)
   int res;
 
   pyobj = check_get_pyobj_ptr(pyptr, NULL);
-  pyobj_key = pycall_pyobject_from_ruby(key);
+  pyobj_key = pycall_convert_index(key);
 
   res = Py_API(PyObject_DelItem)(pyobj, pyobj_key);
   if (res == -1) {
@@ -1731,6 +1760,40 @@ pycall_pydict_from_ruby(VALUE obj)
   return pydictobj;
 }
 
+PyObject *
+pycall_pyslice_from_ruby(VALUE obj)
+{
+  VALUE begin, end, step = Qnil;
+  int exclude_end;
+  PyObject *pystart, *pystop, *pystep, *pyslice;
+
+  if (rb_obj_is_kind_of(obj, rb_cRange)) {
+    pycall_extract_range(obj, &begin, &end, &exclude_end, NULL);
+  }
+  else if (pycall_obj_is_step_range(obj)) {
+    pycall_extract_range(obj, &begin, &end, &exclude_end, &step);
+  }
+  else {
+    rb_raise(rb_eTypeError, "unexpected argument type %s (expected Range or Enumerator generated by Range#step)", rb_class2name(CLASS_OF(obj)));
+  }
+
+  if (!exclude_end) {
+    end = SSIZET2NUM(NUM2SSIZET(end) + 1); /* TODO: limit check */
+  }
+
+  pystart = pycall_pyobject_from_ruby(begin);
+  pystop  = pycall_pyobject_from_ruby(end);
+  pystep  = pycall_pyobject_from_ruby(step);
+
+  pyslice = Py_API(PySlice_New)(pystart, pystop, pystep);
+  /* PySlice_New increments reference counts of pystart, pystop, and pystep */
+  pycall_Py_DecRef(pystart);
+  pycall_Py_DecRef(pystop);
+  pycall_Py_DecRef(pystep);
+
+  return pyslice;
+}
+
 VALUE
 pycall_pyerror_new(PyObject *type, PyObject *value, PyObject *traceback)
 {
diff --git a/ext/pycall/pycall_internal.h b/ext/pycall/pycall_internal.h
@@ -577,6 +577,8 @@ typedef struct {
   PyObject * (* PyTuple_GetItem)(PyObject *, Py_ssize_t);
   int (* PyTuple_SetItem)(PyObject *, Py_ssize_t, PyObject *);
 
+  PyObject * (* PySlice_New)(PyObject *, PyObject *, PyObject *);
+
   PyObject * (* PyIter_Next)(PyObject *);
 
   PyObject * (* PyErr_Occurred)(void);
@@ -653,6 +655,7 @@ PyObject *pycall_pystring_from_ruby(VALUE);
 PyObject *pycall_pytuple_from_ruby(VALUE);
 PyObject *pycall_pylist_from_ruby(VALUE);
 PyObject *pycall_pydict_from_ruby(VALUE);
+PyObject *pycall_pyslice_from_ruby(VALUE);
 
 NORETURN(void pycall_pyerror_fetch_and_raise(char const *format, ...));
 
@@ -662,6 +665,9 @@ PyObject *pycall_pystring_from_formatv(char const *format, va_list vargs);
 
 VALUE pycall_pyrubyptr_new(PyObject *pyrubyobj);
 
+int pycall_obj_is_step_range(VALUE obj);
+int pycall_extract_range(VALUE obj, VALUE *pbegin, VALUE *pend, int *pexclude_end, VALUE *pstep);
+
 void pycall_gcguard_register(PyObject *, VALUE);
 void pycall_gcguard_delete(PyObject *);
 void pycall_gcguard_register_pyrubyobj(PyObject *);
diff --git a/ext/pycall/range.c b/ext/pycall/range.c
@@ -0,0 +1,69 @@
+#include "pycall_internal.h"
+
+struct enumerator_head {
+    VALUE obj;
+    ID    meth;
+    VALUE args;
+};
+
+int
+pycall_obj_is_step_range(VALUE obj)
+{
+  struct enumerator_head *eh;
+
+  if (!RB_TYPE_P(obj, T_DATA)) {
+    return 0;
+  }
+
+  if (!rb_obj_is_kind_of(obj, rb_cEnumerator)) {
+    return 0;
+  }
+
+  eh = (struct enumerator_head *)DATA_PTR(obj);
+
+  if (!rb_obj_is_kind_of(eh->obj, rb_cRange)) {
+    return 0;
+  }
+  if (eh->meth == rb_intern("step")) {
+    if (!RB_TYPE_P(eh->args, T_ARRAY)) {
+      return 0;
+    }
+    return (RARRAY_LEN(eh->args) == 1);
+  }
+
+  return 0;
+}
+
+int
+pycall_extract_range(VALUE obj, VALUE *pbegin, VALUE *pend, int *pexclude_end, VALUE *pstep)
+{
+  ID id_begin, id_end, id_exclude_end;
+  VALUE begin, end, exclude_end, step = Qnil;
+
+  CONST_ID(id_begin,       "begin");
+  CONST_ID(id_end,         "end");
+  CONST_ID(id_exclude_end, "exclude_end?");
+
+  if (rb_obj_is_kind_of(obj, rb_cRange)) {
+    begin = rb_funcallv(obj, id_begin, 0, 0);
+    end   = rb_funcallv(obj, id_end,   0, 0);
+    exclude_end = rb_funcallv(obj, id_exclude_end, 0, 0);
+  }
+  else if (pycall_obj_is_step_range(obj)) {
+    struct enumerator_head *eh = (struct enumerator_head *)DATA_PTR(obj);
+    begin = rb_funcallv(eh->obj, id_begin, 0, 0);
+    end   = rb_funcallv(eh->obj, id_end,   0, 0);
+    exclude_end = rb_funcallv(eh->obj, id_exclude_end, 0, 0);
+    step  = RARRAY_AREF(eh->args, 0);
+  }
+  else {
+    return 0;
+  }
+
+  if (pbegin)       *pbegin = begin;
+  if (pend)         *pend = end;
+  if (pexclude_end) *pexclude_end = RTEST(exclude_end);
+  if (pstep)        *pstep = step;
+
+  return 1;
+}
diff --git a/lib/pycall/pyobject_wrapper.rb b/lib/pycall/pyobject_wrapper.rb
@@ -76,20 +76,10 @@ def #{op}(other)
     end
 
     def [](*key)
-      if key.length == 1
-        key = key[0]
-      else
-        keys = PyCall::Tuple.new(key)
-      end
       LibPython::Helpers.getitem(__pyptr__, key)
     end
 
     def []=(*key, value)
-      if key.length == 1
-        key = key[0]
-      else
-        key = PyCall::Tuple.new(key)
-      end
       LibPython::Helpers.setitem(__pyptr__, key, value)
     end
 
diff --git a/spec/pycall/pyobject_wrapper_spec.rb b/spec/pycall/pyobject_wrapper_spec.rb
@@ -143,11 +143,46 @@ module PyCall
     end
 
     describe '#[]' do
-      pending
+      context 'when the given index is a Range' do
+        specify do
+          list = PyCall::List.new([*1..10])
+          expect(list[1..-2]).to eq(PyCall::List.new([*2..9]))
+          expect(list[1...-2]).to eq(PyCall::List.new([*2..8]))
+        end
+      end
+
+      context 'when the given index is an Enumerable that is created by Range#step' do
+        specify do
+          list = PyCall::List.new([*1..10])
+          expect(list[(1..-2).step(2)]).to eq(PyCall::List.new([2, 4, 6, 8]))
+        end
+      end
     end
 
     describe '#[]=' do
-      pending
+      context 'when the given index is a Range' do
+        specify do
+          list = PyCall::List.new([*1..10])
+          list[1..-2] = [100, 200, 300]
+          expect(list).to eq(PyCall::List.new([1, 100, 200, 300, 10]))
+
+          list = PyCall::List.new([*1..10])
+          list[1...-2] = [100, 200, 300]
+          expect(list).to eq(PyCall::List.new([1, 100, 200, 300, 9, 10]))
+        end
+      end
+
+      context 'when the given index is an Enumerable that is created by Range#step' do
+        specify do
+          list = PyCall::List.new([*1..10])
+          list[(1..-3).step(2)] = [100, 200, 300, 400]
+          expect(list).to eq(PyCall::List.new([1, 100, 3, 200, 5, 300, 7, 400, 9, 10]))
+
+          list = PyCall::List.new([*1..10])
+          list[(1...-3).step(2)] = [100, 200, 300]
+          expect(list).to eq(PyCall::List.new([1, 100, 3, 200, 5, 300, 7, 8, 9, 10]))
+        end
+      end
     end
 
     describe '#call' do
