Respect type overrides in E721 (#3582)

This commit is contained in:
Charlie Marsh 2023-03-17 14:29:05 -04:00 committed by GitHub
parent dedf4cbdeb
commit 2e21920adf
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 24 additions and 16 deletions

View file

@ -54,3 +54,7 @@ if type(a) != type(b) or type(a) == type(ccc):
pass pass
assert type(res) == type(None) assert type(res) == type(None)
types = StrEnum
if x == types.X:
pass

View file

@ -3186,11 +3186,7 @@ where
} }
if self.settings.rules.enabled(Rule::TypeComparison) { if self.settings.rules.enabled(Rule::TypeComparison) {
self.diagnostics.extend(pycodestyle::rules::type_comparison( pycodestyle::rules::type_comparison(self, expr, ops, comparators);
ops,
comparators,
Range::from(expr),
));
} }
if self.settings.rules.enabled(Rule::SysVersionCmpStr3) if self.settings.rules.enabled(Rule::SysVersionCmpStr3)

View file

@ -1,6 +1,7 @@
use itertools::izip; use itertools::izip;
use rustpython_parser::ast::{Cmpop, Constant, Expr, ExprKind}; use rustpython_parser::ast::{Cmpop, Constant, Expr, ExprKind};
use crate::checkers::ast::Checker;
use ruff_diagnostics::{Diagnostic, Violation}; use ruff_diagnostics::{Diagnostic, Violation};
use ruff_macros::{derive_message_formats, violation}; use ruff_macros::{derive_message_formats, violation};
use ruff_python_ast::types::Range; use ruff_python_ast::types::Range;
@ -34,9 +35,7 @@ impl Violation for TypeComparison {
} }
/// E721 /// E721
pub fn type_comparison(ops: &[Cmpop], comparators: &[Expr], location: Range) -> Vec<Diagnostic> { pub fn type_comparison(checker: &mut Checker, expr: &Expr, ops: &[Cmpop], comparators: &[Expr]) {
let mut diagnostics: Vec<Diagnostic> = vec![];
for (op, right) in izip!(ops, comparators) { for (op, right) in izip!(ops, comparators) {
if !matches!(op, Cmpop::Is | Cmpop::IsNot | Cmpop::Eq | Cmpop::NotEq) { if !matches!(op, Cmpop::Is | Cmpop::IsNot | Cmpop::Eq | Cmpop::NotEq) {
continue; continue;
@ -44,8 +43,8 @@ pub fn type_comparison(ops: &[Cmpop], comparators: &[Expr], location: Range) ->
match &right.node { match &right.node {
ExprKind::Call { func, args, .. } => { ExprKind::Call { func, args, .. } => {
if let ExprKind::Name { id, .. } = &func.node { if let ExprKind::Name { id, .. } = &func.node {
// Ex) type(False) // Ex) `type(False)`
if id == "type" { if id == "type" && checker.ctx.is_builtin("type") {
if let Some(arg) = args.first() { if let Some(arg) = args.first() {
// Allow comparison for types which are not obvious. // Allow comparison for types which are not obvious.
if !matches!( if !matches!(
@ -56,7 +55,9 @@ pub fn type_comparison(ops: &[Cmpop], comparators: &[Expr], location: Range) ->
kind: None kind: None
} }
) { ) {
diagnostics.push(Diagnostic::new(TypeComparison, location)); checker
.diagnostics
.push(Diagnostic::new(TypeComparison, Range::from(expr)));
} }
} }
} }
@ -64,15 +65,22 @@ pub fn type_comparison(ops: &[Cmpop], comparators: &[Expr], location: Range) ->
} }
ExprKind::Attribute { value, .. } => { ExprKind::Attribute { value, .. } => {
if let ExprKind::Name { id, .. } = &value.node { if let ExprKind::Name { id, .. } = &value.node {
// Ex) types.IntType // Ex) `types.NoneType`
if id == "types" { if id == "types"
diagnostics.push(Diagnostic::new(TypeComparison, location)); && checker
.ctx
.resolve_call_path(value)
.map_or(false, |call_path| {
call_path.first().map_or(false, |module| *module == "types")
})
{
checker
.diagnostics
.push(Diagnostic::new(TypeComparison, Range::from(expr)));
} }
} }
} }
_ => {} _ => {}
} }
} }
diagnostics
} }