def test_partition_source_indented_encoding():
assert partition_source(' # -*- coding: UTF-8 -*-\n') == [
CodePartition(
CodeType.PRE_IMPORT_CODE,
' # -*- coding: UTF-8 -*-\n',
),
]
def test_partition_source_comment_lines():
assert partition_source('# hello world\n'
'import os\n', ) == [
CodePartition(CodeType.PRE_IMPORT_CODE,
'# hello world\n'),
CodePartition(CodeType.IMPORT, 'import os\n'),
]
def test_partition_source_encoding_and_shebang():
assert partition_source(
'#!/usr/bin/env python\n'
'# -*- coding: UTF-8 -*-\n', ) == [
CodePartition(CodeType.PRE_IMPORT_CODE, '#!/usr/bin/env python\n'),
CodePartition(CodeType.PRE_IMPORT_CODE,
'# -*- coding: UTF-8 -*-\n'),
]
def test_partition_source_unicode_docstring():
assert partition_source('# -*- coding: UTF-8 -*-\n'
'u"""☃☃☃"""\n', ) == [
CodePartition(CodeType.PRE_IMPORT_CODE,
'# -*- coding: UTF-8 -*-\n'),
CodePartition(CodeType.PRE_IMPORT_CODE,
'u"""☃☃☃"""\n'),
]
def test_partition_source_multiple_docstrings():
assert partition_source(
'"""foo"""\n'
'"""bar"""\n', ) == [
# only the first docstring should count as a docstring
CodePartition(CodeType.PRE_IMPORT_CODE, '"""foo"""\n'),
CodePartition(CodeType.PRE_IMPORT_CODE, '"""bar"""\n'),
]
def test_partition_source_blank_lines_with_whitespace():
assert partition_source('import os\n'
'\n'
' \n'
'import sys\n', ) == [
CodePartition(CodeType.IMPORT, 'import os\n'),
CodePartition(CodeType.NON_CODE, '\n'),
CodePartition(CodeType.NON_CODE, ' \n'),
CodePartition(CodeType.IMPORT, 'import sys\n'),
]
def test_partition_source_import_inside_code_not_an_import():
assert partition_source('x = 1\nimport os\n') == [
CodePartition(CodeType.CODE, 'x = 1\nimport os\n'),
]
def test_partition_source_import_contains_comment():
# We want to maintain comments with imports
assert partition_source('from foo import * # noqa\n') == [
CodePartition(CodeType.IMPORT, 'from foo import * # noqa\n'),
]
def test_partion_source_import_no_nl():
assert partition_source('import os') == [
CodePartition(CodeType.IMPORT, 'import os'),
]
def test_partition_source_encoding_no_nl():
assert partition_source('# -*- coding: UTF-8 -*-') == [
CodePartition(CodeType.PRE_IMPORT_CODE, '# -*- coding: UTF-8 -*-'),
]
def test_partition_source_shebang_no_nl():
assert partition_source('#!/usr/bin/env python') == [
CodePartition(CodeType.PRE_IMPORT_CODE, '#!/usr/bin/env python'),
]
def test_partition_source_errors_with_bytes():
with pytest.raises((AttributeError, TypeError)):
partition_source(b'') # type: ignore
def test_partition_source_trivial():
assert partition_source('') == []
def test_partition_source_code_no_nl():
assert partition_source('x = 1') == [
CodePartition(CodeType.CODE, 'x = 1'),
]
def test_partition_source_docstring_no_nl():
assert partition_source('"""foo"""') == [
CodePartition(CodeType.PRE_IMPORT_CODE, '"""foo"""'),
]
def test_partition_source_errors_with_bytes():
with pytest.raises(TypeError):
partition_source(b'')
