saulshanabrook · Nov 1, 2020 · Oct 27, 2020 · Oct 27, 2020 · Oct 27, 2020 · Oct 28, 2020
diff --git a/k8/images/koalas/Dockerfile b/k8/images/koalas/Dockerfile
@@ -26,10 +26,12 @@ RUN --mount=type=cache,target=/root/.cache/pip \
     scikit-learn \
     openpyxl \
     xlrd \
+    pytest-custom_exit_code \
     'plotly>=4.8' && \
     echo $(python -c "import jdk; print(jdk.install('8'))") > /tmp/java_home
 
+ENV PYTHON_RECORD_API_FROM_MODULES=databricks
+
 RUN python -c "import databricks.koalas"
 
-ENV PYTHON_RECORD_API_FROM_MODULES=koalas
-CMD bash -c "JAVA_HOME=$(cat /tmp/java_home) SPARK_LOCAL_IP=127.0.0.1 ./dev/pytest --doctest-modules databricks -k 'not to_clipboard and not test_day_name and not test_month_name'"
+CMD bash -c "JAVA_HOME=$(cat /tmp/java_home) SPARK_LOCAL_IP=127.0.0.1 ./dev/pytest --doctest-modules databricks --suppress-tests-failed-exit-code -k 'not to_clipboard and not test_day_name and not test_month_name'"
diff --git a/k8/images/koalas/version b/k8/images/koalas/version
@@ -1 +1 @@
-1
+3
diff --git a/k8/images/scipy/Dockerfile b/k8/images/scipy/Dockerfile
@@ -18,7 +18,7 @@ RUN git clone \
     .
 
 RUN --mount=type=cache,target=/root/.cache/pip \
-    pip install pytest-custom_exit_code cython numpy pytest
+    pip install pytest-custom_exit_code cython numpy pytest==6.0.1
 
 RUN --mount=type=cache,target=/root/.cache/pip \
     python setup.py build_ext --inplace

diff --git a/k8/images/scipy/version b/k8/images/scipy/version
@@ -1 +1 @@
-0
+1
diff --git a/record_api/__init__.py b/record_api/__init__.py
@@ -4,4 +4,4 @@
 
 from .core import *
 
-__version__ = "1.2.1"
+__version__ = "1.3.2"
diff --git a/record_api/apis.py b/record_api/apis.py
@@ -374,7 +374,6 @@ class Signature(BaseModel):
     var_kw: typing.Optional[typing.Tuple[str, Type]] = None
 
     metadata: typing.Dict[str, int] = pydantic.Field(default_factory=dict)
-    return_type: OutputType = None
 
     @pydantic.validator("pos_only_required")
     @classmethod
@@ -409,13 +408,6 @@ def validate_keys_unique(cls, values) -> None:
             raise ValueError(repr(all_keys))
         return values
 
-    @property
-    def return_type_annotation(self) -> typing.Optional[cst.Annotation]:
-        return_type_annotation = None
-        if self.return_type:
-            return_type_annotation = cst.Annotation(self.return_type.annotation)
-        return return_type_annotation
-
     def function_def(
         self,
         name: str,
@@ -435,7 +427,6 @@ def function_def(
                 [cst.SimpleStatementLine([s]) for s in self.body(indent)]
             ),
             decorators,
-            self.return_type_annotation
         )
 
     def body(self, indent: int) -> typing.Iterable[cst.BaseSmallStatement]:
@@ -517,17 +508,13 @@ def initial_args(self) -> typing.Iterator[Type]:
 
     @classmethod
     def from_params(
-        cls,
-        args: typing.List[object] = [],
-        kwargs: typing.Dict[str, object] = {},
-        return_type: typing.Optional[typing.Dict[str, typing.Union[str, typing.Dict]]] = {},
+        cls, args: typing.List[object] = [], kwargs: typing.Dict[str, object] = {}
     ) -> Signature:
         # If we don't know what the args/kwargs are, assume the args are positional only
         # and the kwargs and keyword only
         return cls(
             pos_only_required={f"_{i}": create_type(v) for i, v in enumerate(args)},
             kw_only_required={k: create_type(v) for k, v in kwargs.items()},
-            return_type=create_type(return_type) if return_type else None
         )
 
     @classmethod
@@ -538,7 +525,6 @@ def from_bound_params(
         var_pos: typing.Optional[typing.Tuple[str, typing.List[object]]] = None,
         kw_only: typing.Dict[str, object] = {},
         var_kw: typing.Optional[typing.Tuple[str, typing.Dict[str, object]]] = None,
-        return_type: typing.Optional[typing.Dict[str, typing.Union[str, typing.Dict]]] = {},
     ) -> Signature:
         return cls(
             pos_only_required={k: create_type(v) for k, v in pos_only},
@@ -552,7 +538,6 @@ def from_bound_params(
                 if var_kw
                 else None
             ),
-            return_type=create_type(return_type) if return_type else None
         )
 
     def content_equal(self, other: Signature) -> bool:
@@ -572,7 +557,6 @@ def __ior__(self, other: Signature) -> Signature:
         self._copy_var_pos(other)
         self._copy_kw_only(other)
         self._copy_var_kw(other)
-        self._copy_return_type(other)
 
         update_add(self.metadata, other.metadata)
         self._trim_positional_only_args()
@@ -734,13 +718,6 @@ def _copy_var_kw(self, other: Signature) -> None:
             else (self.var_kw or other.var_kw)
         )
 
-    def _copy_return_type(self, other: Signature) -> None:
-        self.return_type = (
-            unify((self.return_type, other.return_type,))
-            if self.return_type and other.return_type
-            else (self.return_type or other.return_type)
-        )
-
 
 API.update_forward_refs()
 Class.update_forward_refs()

diff --git a/record_api/core.py b/record_api/core.py
@@ -29,15 +29,6 @@
 context_manager: Optional[ContextManager] = None
 write_line: Optional[Callable[[dict], None]] = None
 
-FUNCTION_CALL_OP_NAMES = {
-    "CALL_METHOD",
-    "CALL_FUNCTION",
-    "CALL_FUNCTION_KW",
-    "CALL_FUNCTION_EX",
-    "LOAD_ATTR",
-    "BINARY_SUBSCR",
-}
-
 
 def get_tracer() -> Tracer:
     global TRACER
@@ -282,7 +273,6 @@ def log_call(
     fn: Callable,
     args: Iterable = (),
     kwargs: Mapping[str, Any] = {},
-    return_type: Any = None,
 ) -> None:
     bound = Bound.create(fn, args, kwargs)
     line: Dict = {"location": location, "function": preprocess(fn)}
@@ -294,8 +284,6 @@ def log_call(
             line["params"]["kwargs"] = {k: preprocess(v) for k, v in kwargs.items()}
     else:
         line["bound_params"] = bound.as_dict()
-    if return_type:
-        line['return_type'] = return_type
     assert write_line
     write_line(line)
 
@@ -307,16 +295,11 @@ class Stack:
     NULL: ClassVar[object] = object()
     current_i: int = dataclasses.field(init=False, default=0)
     opcode: int = dataclasses.field(init=False)
-    previous_stack: Optional[Stack] = None
-    log_call_args: Tuple = ()
 
     def __post_init__(self):
         self.op_stack = get_stack.OpStack(self.frame)
         self.opcode = self.frame.f_code.co_code[self.frame.f_lasti]
 
-        if self.previous_stack and self.previous_stack.previous_stack:
-            del self.previous_stack.previous_stack
-
     @property
     def oparg(self):
         # sort of replicates logic in dis._unpack_opargs but doesn't account for extended
@@ -377,24 +360,14 @@ def pop_n(self, n: int) -> List:
         return l
 
     def process(
-        self,
-        keyed_args: Tuple,
-        fn: Callable,
-        args: Iterable,
-        kwargs: Mapping = {},
-        delay: bool = False
+        self, keyed_args: Tuple, fn: Callable, args: Iterable, kwargs: Mapping = {}
     ) -> None:
-
-        # Note: This take args as an iterable, instead of as a varargs, so that if
-        # we don't trace we don't have to expand the iterable
+        # Note: This take args as an iterable, instead of as a varargs, so that if we don't trace we don't have to expand the iterable
         if self.tracer.should_trace(*keyed_args):
             filename = self.frame.f_code.co_filename
             line = self.frame.f_lineno
             # Don't pass kwargs if not used, so we can more easily test mock calls
-            if not delay:
-                log_call(f"{filename}:{line}", fn, tuple(args), *((kwargs,) if kwargs else ()))
-            else:
-                self.log_call_args = (filename, line, fn, tuple(args), kwargs)
+            log_call(f"{filename}:{line}", fn, tuple(args), *((kwargs,) if kwargs else ()))
 
     def __call__(self) -> None:
         """
@@ -410,34 +383,14 @@ def __call__(self) -> None:
                 (self.TOS, self.TOS1), BINARY_OPS[opname], (self.TOS1, self.TOS)
             )
 
-        if self.previous_stack and self.previous_stack.opname in FUNCTION_CALL_OP_NAMES:
-            self.log_called_method()
-
         method_name = f"op_{opname}"
         if hasattr(self, method_name):
             getattr(self, method_name)()
         return None
 
-    def log_called_method(self):
-        if self.previous_stack.log_call_args:
-            tos = self.TOS
-            if type(tos) is type and issubclass(tos, Exception):
-                # Don't record exception
-                return
-            return_type = type(tos) if type(tos) != type else tos
-            filename, line, fn, args, *kwargs = self.previous_stack.log_call_args
-            kwargs = kwargs[0] if kwargs else {}
-            log_call(
-                f"{filename}:{line}",
-                fn,
-                tuple(args),
-                *((kwargs,) if kwargs else ()),
-                return_type=return_type,
-            )
-
     # special case subscr b/c we only check first arg, not both
     def op_BINARY_SUBSCR(self):
-        self.process((self.TOS1,), op.getitem, (self.TOS1, self.TOS), delay=True)
+        self.process((self.TOS1,), op.getitem, (self.TOS1, self.TOS))
 
     def op_STORE_SUBSCR(self):
         self.process((self.TOS1,), op.setitem, (self.TOS1, self.TOS, self.TOS2))
@@ -446,7 +399,7 @@ def op_DELETE_SUBSCR(self):
         self.process((self.TOS1,), op.delitem, (self.TOS1, self.TOS))
 
     def op_LOAD_ATTR(self):
-        self.process((self.TOS,), getattr, (self.TOS, self.opvalname), delay=True)
+        self.process((self.TOS,), getattr, (self.TOS, self.opvalname))
 
     def op_STORE_ATTR(self):
         self.process((self.TOS,), setattr, (self.TOS, self.opvalname, self.TOS1))
@@ -505,7 +458,7 @@ def op_COMPARE_OP(self):
     def op_CALL_FUNCTION(self):
         args = self.pop_n(self.oparg)
         fn = self.pop()
-        self.process((fn,), fn, args, delay=True)
+        self.process((fn,), fn, args)
 
     def op_CALL_FUNCTION_KW(self):
         kwargs_keys = self.pop()
@@ -515,7 +468,7 @@ def op_CALL_FUNCTION_KW(self):
         args = self.pop_n(self.oparg - n_kwargs)
         fn = self.pop()
 
-        self.process((fn,), fn, args, kwargs, delay=True)
+        self.process((fn,), fn, args, kwargs)
 
     def op_CALL_FUNCTION_EX(self):
         has_kwarg = self.oparg & int("01", 2)
@@ -529,21 +482,20 @@ def op_CALL_FUNCTION_EX(self):
             fn = self.pop()
         if inspect.isgenerator(args):
             return
-        self.process((fn,), fn, args, kwargs, delay=True)
+        self.process((fn,), fn, args, kwargs)
 
     def op_CALL_METHOD(self):
         args = self.pop_n(self.oparg)
         function_or_self = self.pop()
         null_or_method = self.pop()
         if null_or_method is self.NULL:
             function = function_or_self
-            self.process((function,), function, args, delay=True)
+            self.process((function,), function, args)
         else:
             self_ = function_or_self
             method = null_or_method
             self.process(
                 (self_,), method, itertools.chain((self_,), args),
-                delay=True
             )
 
 
@@ -596,7 +548,6 @@ class Tracer:
     calls_to_modules: List[str]
     # the modules we should trace calls from
     calls_from_modules: List[str]
-    previous_stack: Optional[Stack] = None
 
     def __enter__(self):
         sys.settrace(self)
@@ -626,13 +577,7 @@ def __call__(self, frame, event, arg) -> Optional[Tracer]:
             return None
 
         if self.should_trace_frame(frame):
-            stack = Stack(
-                self,
-                frame,
-                previous_stack=self.previous_stack,
-            )
-            stack()
-            self.previous_stack = stack if stack.log_call_args else None
+            Stack(self, frame)()
         return None
 
     def should_trace_frame(self, frame) -> bool:

diff --git a/record_api/infer_apis.py b/record_api/infer_apis.py
@@ -42,16 +42,12 @@ def __main__():
 
 
 def parse_line(
-    n: int,
-    function: object,
-    params=None,
-    bound_params=None,
-    return_type: typing.Optional[typing.Dict[str, typing.Union[str, typing.Dict]]] = None
+    n: int, function: object, params=None, bound_params=None,
 ) -> typing.Optional[API]:
     if bound_params is not None:
-        signature = Signature.from_bound_params(**bound_params, return_type=return_type)
+        signature = Signature.from_bound_params(**bound_params)
     else:
-        signature = Signature.from_params(**params, return_type=return_type)
+        signature = Signature.from_params(**params)
     signature.metadata[f"usage.{LABEL}"] = n
     return process_function(create_type(function), s=signature)
 
@@ -197,6 +193,8 @@ def _setattr(
     properties = {attr: (metadata, value_type)}
     if isinstance(instance, OtherOutput):
         # getting an attribute on a instance
+        if not instance.type.module:
+            return None
         return API(
             modules={
                 instance.type.module: Module(

diff --git a/record_api/jsonl_test.py b/record_api/jsonl_test.py
@@ -0,0 +1,13 @@
+"""
+Test usage of jsonl writer so we can profile memory consumption
+"""
+
+import tempfile
+from .jsonl import write
+
+N = 10000 * 1000
+print(f"Writing {N} lines to temp file")
+with tempfile.NamedTemporaryFile() as fp:
+    with write(fp.name) as writer:
+        for _ in range(N):
+            writer({"hi": "hello"})
diff --git a/record_api/pytest_plugin.py b/record_api/pytest_plugin.py
@@ -5,7 +5,7 @@
 
 
 @pytest.hookimpl(hookwrapper=True)
-def pytest_pyfunc_call(pyfuncitem):
+def pytest_runtest_call(item):
     with get_tracer():
         yield
Original file line number	Diff line number	Diff line change
Expand Up		@@ -4,4 +4,4 @@

		from .core import *

		__version__ = "1.2.1"
		__version__ = "1.3.2"