From c0b49eab41451ab29ec2c0ba255a66365025b59e Mon Sep 17 00:00:00 2001
From: Robert Bradshaw <robertwb@math.washington.edu>
Date: Fri, 15 Jan 2010 17:35:47 -0800
Subject: [PATCH] Support for [] operator.

---
 Cython/Compiler/ExprNodes.py | 34 +++++++++++++++++++++++++++-------
 Cython/Compiler/Parsing.py   |  3 ++-
 tests/run/cpp_stl_vector.pyx | 17 +++++++++++++++++
 3 files changed, 46 insertions(+), 8 deletions(-)

diff --git a/Cython/Compiler/ExprNodes.py b/Cython/Compiler/ExprNodes.py
index afbf2f8e..6e9388b6 100755
--- a/Cython/Compiler/ExprNodes.py
+++ b/Cython/Compiler/ExprNodes.py
@@ -1914,18 +1914,37 @@ class IndexNode(ExprNode):
             else:
                 if self.base.type.is_ptr or self.base.type.is_array:
                     self.type = self.base.type.base_type
+                    if self.index.type.is_pyobject:
+                        self.index = self.index.coerce_to(
+                            PyrexTypes.c_py_ssize_t_type, env)
+                    if not self.index.type.is_int:
+                        error(self.pos,
+                            "Invalid index type '%s'" %
+                                self.index.type)
+                elif self.base.type.is_cpp_class:
+                    function = self.base.type.scope.lookup("operator[]")
+                    if function is None:
+                        error(self.pos, "Indexing '%s' not supported" % self.base.type)
+                    else:
+                        function = PyrexTypes.best_match([self.index], function.all_alternatives(), self.pos)
+                        if function is None:
+                            error(self.pos, "Invalid index type '%s'" % self.index.type)
+                    if function is None:
+                        self.type = PyrexTypes.error_type
+                        self.result_code = "<error>"
+                        return
+                    func_type = function.type
+                    if func_type.is_ptr:
+                        func_type = func_type.base_type
+                    self.index = self.index.coerce_to(func_type.args[0].type, env)
+                    self.type = func_type.return_type
+                    if setting and not func_type.return_type.is_reference:
+                        error(self.pos, "Can't set non-reference '%s'" % self.type)
                 else:
                     error(self.pos,
                         "Attempting to index non-array type '%s'" %
                             self.base.type)
                     self.type = PyrexTypes.error_type
-                if self.index.type.is_pyobject:
-                    self.index = self.index.coerce_to(
-                        PyrexTypes.c_py_ssize_t_type, env)
-                if not self.index.type.is_int:
-                    error(self.pos,
-                        "Invalid index type '%s'" %
-                            self.index.type)
     gil_message = "Indexing Python object"
 
     def nogil_check(self, env):
@@ -4738,6 +4757,7 @@ class NumBinopNode(BinopNode):
         if type2.is_ptr:
             type2 = type2.base_type
         entry = env.lookup(type1.name)
+        # Shouldn't this be type1.scope?
         function = entry.type.scope.lookup("operator%s" % self.operator)
         if function is not None:
             operands = [self.operand2]
diff --git a/Cython/Compiler/Parsing.py b/Cython/Compiler/Parsing.py
index bf638e95..b1dd1847 100644
--- a/Cython/Compiler/Parsing.py
+++ b/Cython/Compiler/Parsing.py
@@ -2061,6 +2061,7 @@ supported_overloaded_operators = set([
     '+', '-', '*', '/', '%', 
     '++', '--', '~', '|', '&', '^', '<<', '>>',
     '==', '!=', '>=', '>', '<=', '<',
+    '[]',
 ])
 
 def p_c_simple_declarator(s, ctx, empty, is_type, cmethod_flag,
@@ -2108,6 +2109,7 @@ def p_c_simple_declarator(s, ctx, empty, is_type, cmethod_flag,
             cname = ctx.namespace + "::" + name
         if name == 'operator' and ctx.visibility == 'extern':
             op = s.sy
+            s.next()
             # Handle diphthong operators.
             if op == '(':
                 s.expect(')')
@@ -2115,7 +2117,6 @@ def p_c_simple_declarator(s, ctx, empty, is_type, cmethod_flag,
             elif op == '[':
                 s.expect(']')
                 op = '[]'
-            s.next()
             if op in ['-', '+', '|', '&'] and s.sy == op:
                 op = op*2
                 s.next()
diff --git a/tests/run/cpp_stl_vector.pyx b/tests/run/cpp_stl_vector.pyx
index cdcb71ba..f6a81b57 100644
--- a/tests/run/cpp_stl_vector.pyx
+++ b/tests/run/cpp_stl_vector.pyx
@@ -3,6 +3,7 @@ cdef extern from "<vector>" namespace std:
     cdef cppclass vector[T]:
         void push_back(T)
         size_t size()
+        T operator[](size_t)
 
 def simple_test(double x):
     """
@@ -37,3 +38,19 @@ def list_test(L):
         return len(L), v.size()
     finally:
         del v
+
+def index_test(L):
+    """
+    >>> index_test([1,2,4,8])
+    (1.0, 8.0)
+    >>> index_test([1.25])
+    (1.25, 1.25)
+    """
+    cdef vector[double] *v 
+    try:
+        v = new vector[double]()
+        for a in L:
+            v.push_back(a)
+        return v[0][0], v[0][len(L)-1]
+    finally:
+        del v
-- 
2.26.2