Use PyObject_CheckReadBuffer().

2 files changed, 4 insertions, 17 deletions

diff --git a/Python/compile.c b/Python/compile.c
index 1a46064..b477513 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -268,7 +268,6 @@ PyCode_New(int argcount, int nlocals, int stacksize, int flags,
 {
 	PyCodeObject *co;
 	int i;
-	PyBufferProcs *pb;
 	/* Check argument types */
 	if (argcount < 0 || nlocals < 0 ||
 	    code == NULL ||
@@ -279,16 +278,8 @@ PyCode_New(int argcount, int nlocals, int stacksize, int flags,
 	    cellvars == NULL || !PyTuple_Check(cellvars) ||
 	    name == NULL || !PyString_Check(name) ||
 	    filename == NULL || !PyString_Check(filename) ||
-		lnotab == NULL || !PyString_Check(lnotab)) {
-		PyErr_BadInternalCall();
-		return NULL;
-	}
-	pb = code->ob_type->tp_as_buffer;
-	if (pb == NULL ||
-	    pb->bf_getreadbuffer == NULL ||
-	    pb->bf_getsegcount == NULL ||
-	    (*pb->bf_getsegcount)(code, NULL) != 1)
-	{
+	    lnotab == NULL || !PyString_Check(lnotab) ||
+	    !PyObject_CheckReadBuffer(code)) {
 		PyErr_BadInternalCall();
 		return NULL;
 	}
diff --git a/Python/marshal.c b/Python/marshal.c
index 806b91c..0d21e84 100644
--- a/Python/marshal.c
+++ b/Python/marshal.c
@@ -108,7 +108,6 @@ static void
 w_object(PyObject *v, WFILE *p)
 {
 	int i, n;
-	PyBufferProcs *pb;
 
 	p->depth++;
 
@@ -249,13 +248,10 @@ w_object(PyObject *v, WFILE *p)
 		w_short(co->co_firstlineno, p);
 		w_object(co->co_lnotab, p);
 	}
-	else if ((pb = v->ob_type->tp_as_buffer) != NULL &&
-		 pb->bf_getsegcount != NULL &&
-		 pb->bf_getreadbuffer != NULL &&
-		 (*pb->bf_getsegcount)(v, NULL) == 1)
-	{
+	else if (PyObject_CheckReadBuffer(v)) {
 		/* Write unknown buffer-style objects as a string */
 		char *s;
+		PyBufferProcs *pb = v->ob_type->tp_as_buffer;
 		w_byte(TYPE_STRING, p);
 		n = (*pb->bf_getreadbuffer)(v, 0, (void **)&s);
 		w_long((long)n, p);