fixed several edge cases when loading lines from files, and added tests

docs/configuration.asciidoc

@@ -293,31 +293,29 @@ With this setting, you can limit the length of lists in local variables.
 [[config-source-lines-error-app-frames]]
 ==== `source_lines_error_app_frames`
 [float]
-[[config-source-lines-span-app-frames]]
-==== `source_lines_span_app_frames`
-[float]
 [[config-source-lines-error-library-frames]]
 ==== `source_lines_error_library_frames`
 [float]
+[[config-source-lines-span-app-frames]]
+==== `source_lines_span_app_frames`
+[float]
 [[config-source-lines-span-library-frames]]
 ==== `source_lines_span_library_frames`
 
 |============
 | Environment                                     | Django/Flask                        | Default
 | `ELASTIC_APM_SOURCE_LINES_ERROR_APP_FRAMES`     | `SOURCE_LINES_ERROR_APP_FRAMES`     | `5`
-| `ELASTIC_APM_SOURCE_LINES_SPAN_APP_FRAMES`      | `SOURCE_LINES_SPAN_APP_FRAMES`      | `5`
 | `ELASTIC_APM_SOURCE_LINES_ERROR_LIBRARY_FRAMES` | `SOURCE_LINES_ERROR_LIBRARY_FRAMES` | `5`
+| `ELASTIC_APM_SOURCE_LINES_SPAN_APP_FRAMES`      | `SOURCE_LINES_SPAN_APP_FRAMES`      | `0`
 | `ELASTIC_APM_SOURCE_LINES_SPAN_LIBRARY_FRAMES`  | `SOURCE_LINES_SPAN_LIBRARY_FRAMES`  | `0`
 |============
 
-By default, the APM agent collects source code snippets for
-
- * errors, both library frames and in-app frames
- * transaction spans, only for in-app frames.
-
+By default, the APM agent collects source code snippets for errors.
 With the above settings, you can modify how many lines of source code is collected.
 
-WARNING: Especially for transactions, collecting source code can have a large impact on storage use in your Elasticsearch cluster.
+We differ between errors and spans, as well as library frames and app frames.
+
+WARNING: Especially for spans, collecting source code can have a large impact on storage use in your Elasticsearch cluster.
 
 [float]
 [[config-flush-interval]]

elasticapm/base.py

@@ -26,8 +26,9 @@
 from elasticapm.conf import Config, constants
 from elasticapm.traces import TransactionsStore, get_transaction
 from elasticapm.transport.base import TransportException
+from elasticapm.utils import compat, is_master_process
 from elasticapm.utils import json_encoder as json
-from elasticapm.utils import compat, is_master_process, stacks, varmap
+from elasticapm.utils import stacks, varmap
 from elasticapm.utils.encoding import shorten, transform
 from elasticapm.utils.module_import import import_string
 

elasticapm/conf/__init__.py

@@ -160,10 +160,10 @@ class Config(_ConfigBase):
     transaction_max_spans = _ConfigValue('TRANSACTION_MAX_SPANS', type=int, default=500)
     max_queue_size = _ConfigValue('MAX_QUEUE_SIZE', type=int, default=500)
     collect_local_variables = _ConfigValue('COLLECT_LOCAL_VARIABLES', default='errors')
+    source_lines_error_app_frames = _ConfigValue('SOURCE_LINES_ERROR_APP_FRAMES', type=int, default=5)
     source_lines_error_library_frames = _ConfigValue('SOURCE_LINES_ERROR_LIBRARY_FRAMES', type=int, default=5)
+    source_lines_span_app_frames = _ConfigValue('SOURCE_LINES_SPAN_APP_FRAMES', type=int, default=0)
     source_lines_span_library_frames = _ConfigValue('SOURCE_LINES_SPAN_LIBRARY_FRAMES', type=int, default=0)
-    source_lines_error_app_frames = _ConfigValue('SOURCE_LINES_ERROR_APP_FRAMES', type=int, default=5)
-    source_lines_span_app_frames = _ConfigValue('SOURCE_LINES_SPAN_APP_FRAMES', type=int, default=5)
     local_var_max_length = _ConfigValue('LOCAL_VAR_MAX_LENGTH', type=int, default=200)
     local_var_list_max_length = _ConfigValue('LOCAL_VAR_LIST_MAX_LENGTH', type=int, default=10)
     async_mode = _BoolConfigValue('ASYNC_MODE', default=True)

elasticapm/contrib/django/__init__.py

@@ -10,5 +10,4 @@
 """
 from elasticapm.contrib.django.client import *  # noqa E401
 
-
 default_app_config = 'elasticapm.contrib.django.apps.ElasticAPMConfig'

elasticapm/utils/stacks.py

@@ -10,6 +10,7 @@
 """
 import fnmatch
 import inspect
+import itertools
 import os
 import re
 import sys
@@ -32,48 +33,56 @@ def get_lines_from_file(filename, lineno, context_lines, loader=None, module_nam
     Returns context_lines before and after lineno from file.
     Returns (pre_context_lineno, pre_context, context_line, post_context).
     """
+    lineno = lineno - 1
+    lower_bound = max(0, lineno - context_lines)
+    upper_bound = lineno + context_lines
+
     source = None
     if loader is not None and hasattr(loader, "get_source"):
-        try:
-            source = loader.get_source(module_name)
-        except ImportError:
-            # ImportError: Loader for module cProfile cannot handle module __main__
-            source = None
-        if source is not None:
-            source = source.splitlines()
+        result = get_source_lines_from_loader(loader, module_name, lineno, lower_bound, upper_bound)
+        if result is not None:
+            return result
+
     if source is None:
         try:
-            f = open(filename, 'rb')
-            try:
-                source = f.readlines()
-            finally:
-                f.close()
-        except (OSError, IOError):
+            with open(filename, 'rb') as file_obj:
+                encoding = 'utf8'
+                # try to find encoding of source file by "coding" header
+                # if none is found, utf8 is used as a fallback
+                for line in itertools.islice(file_obj, 0, 2):
+                    match = _coding_re.search(line.decode('utf8'))
+                    if match:
+                        encoding = match.group(1)
+                        break
+                file_obj.seek(0)
+                lines = [compat.text_type(line, encoding, 'replace')
+                         for line in itertools.islice(file_obj, lower_bound, upper_bound + 1)]
+                offset = lineno - lower_bound
+                return (
+                    [l.strip('\r\n') for l in lines[0:offset]],
+                    lines[offset].strip('\r\n'),
+                    [l.strip('\r\n') for l in lines[offset + 1:]] if len(lines) > offset else []
+                )
+        except (OSError, IOError, IndexError):
             pass
+    return None, None, None
 
-        if source is None:
-            return None, None, None
-        encoding = 'utf8'
-        for line in source[:2]:
-            # File coding may be specified. Match pattern from PEP-263
-            # (http://www.python.org/dev/peps/pep-0263/)
-            match = _coding_re.search(line.decode('utf8'))  # let's assume utf8
-            if match:
-                encoding = match.group(1)
-                break
-        source = [compat.text_type(sline, encoding, 'replace') for sline in source]
-
-    lower_bound = max(0, lineno - context_lines)
-    upper_bound = lineno + 1 + context_lines
 
+def get_source_lines_from_loader(loader, module_name, lineno, lower_bound, upper_bound):
+    try:
+        source = loader.get_source(module_name)
+    except ImportError:
+        # ImportError: Loader for module cProfile cannot handle module __main__
+        return None
+    if source is not None:
+        source = source.splitlines()
     try:
         pre_context = [line.strip('\r\n') for line in source[lower_bound:lineno]]
         context_line = source[lineno].strip('\r\n')
-        post_context = [line.strip('\r\n') for line in source[(lineno + 1):upper_bound]]
+        post_context = [line.strip('\r\n') for line in source[(lineno + 1):upper_bound + 1]]
     except IndexError:
         # the file may have changed since it was loaded into memory
         return None, None, None
-
     return pre_context, context_line, post_context
 
 
@@ -181,9 +190,6 @@ def get_frame_info(frame, lineno, with_locals=True,
         abs_path = None
         function = None
 
-    if lineno:
-        lineno -= 1
-
     # Try to pull a relative file path
     # This changes /foo/site-packages/baz/bar.py into baz/bar.py
     try:
@@ -200,7 +206,7 @@ def get_frame_info(frame, lineno, with_locals=True,
         'filename': filename,
         'module': module_name,
         'function': function,
-        'lineno': lineno + 1,
+        'lineno': lineno,
         'library_frame': is_library_frame(abs_path, include_paths_re, exclude_paths_re)
     }
 
@@ -212,11 +218,9 @@ def get_frame_info(frame, lineno, with_locals=True,
     else:
         pre_context, context_line, post_context = [], None, []
     if context_line:
-        frame_result.update({
-            'pre_context': pre_context,
-            'context_line': context_line,
-            'post_context': post_context,
-        })
+        frame_result['pre_context'] = pre_context
+        frame_result['context_line'] = context_line
+        frame_result['post_context'] = post_context
     if with_locals:
         if f_locals is not None and not isinstance(f_locals, dict):
             # XXX: Genshi (and maybe others) have broken implementations of

tests/utils/json/tests.py

@@ -4,8 +4,8 @@
 import datetime
 import uuid
 
-from elasticapm.utils import json_encoder as json
 from elasticapm.utils import compat
+from elasticapm.utils import json_encoder as json
 
 
 def test_uuid():

tests/utils/stacks/linenos.py

@@ -0,0 +1,20 @@
+1
+2
+3
+4
+5
+6
+7
+8
+9
+10
+11
+12
+13
+14
+15
+16
+17
+18
+19
+20

tests/utils/stacks/tests.py

@@ -2,6 +2,7 @@
 from __future__ import absolute_import
 
 import os
+import pkgutil
 
 import pytest
 from mock import Mock
@@ -113,3 +114,37 @@ def test_get_frame_info():
     assert frame_info['lineno'] == 6
     assert frame_info['context_line'] == '    return inspect.currentframe()'
     assert frame_info['vars'] == {'a_local_var': 42}
+
+
+@pytest.mark.parametrize("lineno,context,expected", [
+    (10, 5, (['5', '6', '7', '8', '9'], '10', ['11', '12', '13', '14', '15'])),
+    (1, 5, ([], '1', ['2', '3', '4', '5', '6'])),
+    (2, 5, (['1'], '2', ['3', '4', '5', '6', '7'])),
+    (20, 5, (['15', '16', '17', '18', '19'], '20', [])),
+    (19, 5, (['14', '15', '16', '17', '18'], '19', ['20'])),
+    (1, 0, ([], '1', [])),
+    (21, 0, (None, None, None)),
+])
+def test_get_lines_from_file(lineno, context, expected):
+    stacks.get_lines_from_file.cache_clear()
+    fname = os.path.join(os.path.dirname(__file__), 'linenos.py')
+    result = stacks.get_lines_from_file(fname, lineno, context)
+    assert result == expected
+
+
+@pytest.mark.parametrize("lineno,context,expected", [
+    (10, 5, (['5', '6', '7', '8', '9'], '10', ['11', '12', '13', '14', '15'])),
+    (1, 5, ([], '1', ['2', '3', '4', '5', '6'])),
+    (2, 5, (['1'], '2', ['3', '4', '5', '6', '7'])),
+    (20, 5, (['15', '16', '17', '18', '19'], '20', [])),
+    (19, 5, (['14', '15', '16', '17', '18'], '19', ['20'])),
+    (1, 0, ([], '1', [])),
+    (21, 0, (None, None, None)),
+])
+def test_get_lines_from_loader(lineno, context, expected):
+    stacks.get_lines_from_file.cache_clear()
+    module = 'tests.utils.stacks.linenos'
+    loader = pkgutil.get_loader(module)
+    fname = os.path.join(os.path.dirname(__file__), 'linenos.py')
+    result = stacks.get_lines_from_file(fname, lineno, context, loader=loader, module_name=module)
+    assert result == expected