# Copyright (c) Microsoft Corporation. All rights reserved.

# Licensed under the MIT License. See LICENSE in the project root

# for license information.

"""Helpers to work with Python code.

"""

import py.path

import re

_marked_line_numbers_cache = {}

def get_marked_line_numbers(path):

"""Given a path to a Python source file, extracts line numbers for all lines

that are marked with #@. For example, given this file::

print(1) # @foo

print(2)

print(3) # @bar,baz

the function will return::

{"foo": 1, "bar": 3, "baz": 3}

"""

if isinstance(path, py.path.local):

path = path.strpath

try:

return _marked_line_numbers_cache[path]

except KeyError:

pass

# Read as bytes to avoid decoding errors.

with open(path, "rb") as f:

lines = {}

for i, line in enumerate(f):

match = re.search(rb"#\s*@(.+?)\s*$", line)

if match:

markers = match.group(1).decode("ascii")

for marker in markers.split(","):

lines[marker] = i + 1

_marked_line_numbers_cache[path] = lines

return lines