Skip to content

Commit

Permalink
Reworked consider-using-str-partition
Browse files Browse the repository at this point in the history
  • Loading branch information
@yushao2
yushao2 committed May 13, 2021
1 parent 943bd58 commit ab21dce
Show file tree
Hide file tree
Showing 3 changed files with 176 additions and 156 deletions.
236 changes: 133 additions & 103 deletions pylint/checkers/refactoring/recommendation_checker.py
View file
@@ -1,6 +1,5 @@
# Licensed under the GPL: https://www.gnu.org/licenses/old-licenses/gpl-2.0.html
# For details: https://github.com/PyCQA/pylint/blob/master/LICENSE
import re
from typing import cast

import astroid
Expand Down Expand Up @@ -37,12 +36,14 @@ class RecommendationChecker(checkers.BaseChecker):
"method of the dictionary instead.",
),
"C0207": (
"Consider using %s[%d] instead",
"Consider using %s instead",
"consider-using-str-partition",
"Emitted when accessing only the first or last element of a str.split(sep). "
"Or when a str.split(sep,maxsplit=1) is used. "
"The first and last element can be accessed by using str.partition(sep)[0] "
"or str.rpartition(sep)[-1] instead, which is less computationally "
"expensive.",
"expensive and works the same as str.split() or str.rsplit() with a maxsplit "
"of 1",
),
}

Expand All @@ -53,8 +54,14 @@ def _is_builtin(node, function):
return False
return utils.is_builtin_object(inferred) and inferred.name == function

@utils.check_messages("consider-iterating-dictionary")
def visit_call(self, node):
@utils.check_messages(
"consider-iterating-dictionary", "consider-using-str-partition"
)
def visit_call(self, node: astroid.Call) -> None:
self._check_consider_iterating_dictionary(node)
self._check_consider_using_str_partition(node)

def _check_consider_iterating_dictionary(self, node: astroid.Call) -> None:
if not isinstance(node.func, astroid.Attribute):
return
if node.func.attrname != "keys":
Expand All @@ -71,6 +78,127 @@ def visit_call(self, node):
if isinstance(node.parent, (astroid.For, astroid.Comprehension)):
self.add_message("consider-iterating-dictionary", node=node)

def _check_consider_using_str_partition(self, node: astroid.Call) -> None:
"""Add message when accessing first or last elements of a str.split() or str.rsplit(),
or when split/rsplit with max_split=1 is used"""

# Check if call is split() or rsplit()
if isinstance(node.func, astroid.Attribute) and node.func.attrname in (
"split",
"rsplit",
):
inferred_func = utils.safe_infer(node.func)
fn_name = node.func.attrname
node_name = node.as_string()

if not isinstance(inferred_func, astroid.BoundMethod):
return

try:
seperator = utils.get_argument_from_call(node, 0, "sep").value
except utils.NoSuchArgumentError:
return
# Check if maxsplit is set, and if it's == 1 (this can be replaced by partition)
try:
maxsplit = utils.get_argument_from_call(node, 1, "maxsplit").value
if maxsplit == 1:
self.add_message(
"consider-using-str-partition",
node=node,
args=(
f"{node.func.expr.as_string()}.{node.func.attrname.replace('split','partition')}('{seperator}')"
),
)
return

except utils.NoSuchArgumentError:
pass

# Check if it's immediately subscripted
if isinstance(node.parent, astroid.Subscript):
subscript_node = node.parent
# Check if subscripted with -1/0
if not isinstance(
subscript_node.slice, (astroid.Const, astroid.UnaryOp)
):
return
subscript_value = utils.safe_infer(subscript_node.slice)
subscript_value = cast(astroid.Const, subscript_value)
subscript_value = subscript_value.value
if subscript_value in (-1, 0):
new_fn = "rpartition" if subscript_value == -1 else "partition"
new_fn = new_fn[::-1]
node_name = node_name[::-1]
fn_name = fn_name[::-1]
new_name = node_name.replace(fn_name, new_fn, 1)[::-1]
self.add_message(
"consider-using-str-partition", node=node, args=(new_name,)
)
return

# Check where name it's assigned to, then check all usage of name
if isinstance(node.parent, astroid.Tuple):
assign_node = node.parent.parent
if not isinstance(assign_node, astroid.Assign):
return
idx = node.parent.elts.index(node)
if not isinstance(assign_node.targets[0], astroid.Tuple):
return
assign_target = assign_node.targets[0].elts[idx]
elif isinstance(node.parent, astroid.Assign):
assign_target = node.parent.targets[0]
else:
return

if not isinstance(assign_target, astroid.AssignName):
return

# Go to outer-most scope (module), then search the child for usage
module_node = node
while not isinstance(module_node, astroid.Module):
module_node = module_node.parent

subscript_usage = set()
for child in module_node.body:
for search_node in child.nodes_of_class(astroid.Name):
search_node = cast(astroid.Name, search_node)

last_definition = search_node.lookup(search_node.name)[1][-1]
if last_definition is not assign_target:
continue
if not isinstance(search_node.parent, astroid.Subscript):
continue
subscript_node = search_node.parent
if not isinstance(
subscript_node.slice, (astroid.Const, astroid.UnaryOp)
):
return
subscript_value = utils.safe_infer(subscript_node.slice)
subscript_value = cast(astroid.Const, subscript_value)
subscript_value = subscript_value.value
if subscript_value not in (-1, 0):
return
subscript_usage.add(subscript_value)

# Construct help text
help_text = ""
node_name = node_name[::-1]
fn_name = fn_name[::-1]
if 0 in subscript_usage:
new_fn = "partition"[::-1]
new_name = node_name.replace(fn_name, new_fn, 1)[::-1]
help_text += f"{new_name} to extract the first element of a .split()"

if -1 in subscript_usage:
new_fn = "rpartition"[::-1]
new_name = node_name.replace(fn_name, new_fn, 1)[::-1]
help_text += " and " if help_text != "" else ""
help_text += f"{new_name} to extract the last element of a .split()"

self.add_message(
"consider-using-str-partition", node=node, args=(help_text,)
)

@utils.check_messages("consider-using-enumerate", "consider-using-dict-items")
def visit_for(self, node: astroid.For) -> None:
self._check_consider_using_enumerate(node)
Expand Down Expand Up @@ -219,101 +347,3 @@ def visit_comprehension(self, node: astroid.Comprehension) -> None:

self.add_message("consider-using-dict-items", node=node)
return

@utils.check_messages("consider-using-str-partition")
def visit_subscript(self, node: astroid.Subscript) -> None:
"""Add message when accessing first or last elements of a str.split()."""
assignment = None
subscripted_object = node.value
if isinstance(subscripted_object, astroid.Name):
# Check assignment of this subscripted object
assignment_lookup_results = node.value.lookup(node.value.name)
if (
len(assignment_lookup_results) == 0
or len(assignment_lookup_results[-1]) == 0
):
return
assign_name = assignment_lookup_results[-1][-1]
if not isinstance(assign_name, astroid.AssignName):
return
assignment = assign_name.parent
if not isinstance(assignment, astroid.Assign):
return
# Set subscripted_object to the assignment RHS
subscripted_object = assignment.value

if isinstance(subscripted_object, astroid.Call):
# Check if call is .split() or .rsplit()
if isinstance(
subscripted_object.func, astroid.Attribute
) and subscripted_object.func.attrname in ["split", "rsplit"]:
inferred = utils.safe_infer(subscripted_object.func)
if not isinstance(inferred, astroid.BoundMethod):
return

# Check if subscript is 1 or -1
value = node.slice
if isinstance(value, astroid.Index):
value = value.value

if isinstance(value, (astroid.UnaryOp, astroid.Const)):
const = utils.safe_infer(value)
const = cast(astroid.Const, const)
p = re.compile(
r"tilpsr*\."
) # Match and replace in reverse to ensure last occurance replaced
if const.value == -1:
help_text = p.sub(
".rpartition"[::-1],
subscripted_object.as_string()[::-1],
count=1,
)[::-1]
self.add_message(
"consider-using-str-partition",
node=node,
args=(help_text, -1),
)
elif const.value == 0:
help_text = p.sub(
".partition"[::-1],
subscripted_object.as_string()[::-1],
count=1,
)[::-1]
self.add_message(
"consider-using-str-partition",
node=node,
args=(help_text, 0),
)
# Check if subscript is len(split) - 1
if (
isinstance(value, astroid.BinOp)
and value.left.func.name == "len"
and value.op == "-"
and value.right.value == 1
and assignment is not None
):
# Ensure that the len() is from builtins, not user re-defined
inferred_fn = utils.safe_infer(value.left.func)
if not (
isinstance(inferred_fn, astroid.FunctionDef)
and isinstance(inferred_fn.parent, astroid.Module)
and inferred_fn.parent.name == "builtins"
):
return

# Further check for argument within len(), and compare that to object being split
arg_within_len = value.left.args[0].name
if arg_within_len == node.value.name:
p = re.compile(
r"tilpsr*\."
) # Match and replace in reverse to ensure last occurance replaced
help_text = p.sub(
".rpartition"[::-1],
subscripted_object.as_string()[::-1],
count=1,
)[::-1]
self.add_message(
"consider-using-str-partition",
node=node,
args=(help_text, -1),
)
42 changes: 18 additions & 24 deletions tests/functional/c/consider/consider_using_str_partition.py
View file
@@ -1,7 +1,8 @@
"""Emit a message for iteration through dict keys and subscripting dict with key."""
# pylint: disable=line-too-long,missing-docstring,unsubscriptable-object,too-few-public-methods,invalid-name,redefined-builtin
from collections.abc import Iterable
from typing import Any

get_first = '1,2,3'.split(sep=',',maxsplit=1)[0] # [consider-using-str-partition]
get_last = '1,2,3'[::-1].split(',',1)[0] # [consider-using-str-partition]

get_first = '1,2,3'.split(',')[0] # [consider-using-str-partition]
get_last = '1,2,3'[::-1].split(',')[0] # [consider-using-str-partition]
Expand All @@ -16,21 +17,21 @@
get_mid = SEQ.split(',')[-2] # This is okay

# Test with storing splits as an object
split1 = SEQ.split(',')
get_first = split1[0] # [consider-using-str-partition]
get_last = split1[-1] # [consider-using-str-partition]
get_last = split1[len(split1)-1] # [consider-using-str-partition]
split1 = SEQ.split(',') # [consider-using-str-partition]
get_first = split1[0]
get_last = split1[-1]

split2 = SEQ.rsplit(',')
get_first = split2[0] # [consider-using-str-partition]
get_last = split2[-1] # [consider-using-str-partition]
get_last = split2[len(split2)-1] # [consider-using-str-partition]
split2 = SEQ.rsplit(',') # [consider-using-str-partition]
get_first = split2[0]
get_last = split2[-1]

print(split1[0], split1[-1]) # [consider-using-str-partition, consider-using-str-partition]
split1, split2 = SEQ.split(','), SEQ.rsplit(',') # [consider-using-str-partition, consider-using-str-partition]
get_first = split1[0]
get_last = split1[-1]
get_first = split2[0]
get_last = split2[-1]

# Test when running len on another iterable
some_list = []
get_last = split1[len(split2)-1] # Should not throw an error
print(split1[0], split1[-1])

# Tests on class attributes
class Foo():
Expand All @@ -51,10 +52,9 @@ def get_string(self) -> str:
get_mid = Foo.class_str.split(',')[1]
get_mid = Foo.class_str.split(',')[-2]

split2 = Foo.class_str.split(',')
get_first = split2[0] # [consider-using-str-partition]
get_last = split2[-1] # [consider-using-str-partition]
get_last = split2[len(split2)-1] # [consider-using-str-partition]
split2 = Foo.class_str.split(',') # [consider-using-str-partition]
get_first = split2[0]
get_last = split2[-1]

# Test with accessors
bar = Foo()
Expand All @@ -79,9 +79,3 @@ class Bar():
print(Bar.split.split(",")[-1]) # [consider-using-str-partition] (Error message should show Bar.split.rpartition)
print(Bar.split.rsplit(",")[0]) # [consider-using-str-partition] (Error message should show Bar.split.partition)
print(Bar.split.rsplit(",")[-1]) # [consider-using-str-partition] (Error message should show Bar.split.rpartition)

# Test with user-defined len function
def len(x: Iterable[Any]) -> str:
return f"Hello, world! {x[2]}"

get_last = split2[len(split2)-1] # This won't throw the warning as the len function has been redefined

0 comments on commit ab21dce

Please sign in to comment.