From 86806a9e39ed0ace7bbad87db62ca2d22ec3e5cf Mon Sep 17 00:00:00 2001 From: Martin Lehoux Date: Sat, 19 Nov 2022 15:29:05 +0100 Subject: [PATCH] U013: Also convert typing.TypedDict (#810) --- resources/test/fixtures/U013.py | 4 + .../convert_typed_dict_functional_to_class.rs | 53 ++++++----- .../ruff__linter__tests__U013_U013.py.snap | 89 +++++++++++-------- 3 files changed, 90 insertions(+), 56 deletions(-) diff --git a/resources/test/fixtures/U013.py b/resources/test/fixtures/U013.py index 663079d5cd..28c983017f 100644 --- a/resources/test/fixtures/U013.py +++ b/resources/test/fixtures/U013.py @@ -1,4 +1,5 @@ from typing import TypedDict, NotRequired, Literal +import typing # dict literal MyType1 = TypedDict("MyType1", {"a": int, "b": str}) @@ -27,3 +28,6 @@ MyType9 = TypedDict("MyType9", {"in": int, "x-y": int}) # using Literal type MyType10 = TypedDict("MyType10", {"key": Literal["value"]}) + +# using namespace TypedDict +MyType11 = typing.TypedDict("MyType11", {"key": int}) diff --git a/src/pyupgrade/plugins/convert_typed_dict_functional_to_class.rs b/src/pyupgrade/plugins/convert_typed_dict_functional_to_class.rs index 5625d244ee..e5cdc0a64c 100644 --- a/src/pyupgrade/plugins/convert_typed_dict_functional_to_class.rs +++ b/src/pyupgrade/plugins/convert_typed_dict_functional_to_class.rs @@ -2,6 +2,7 @@ use anyhow::{bail, Result}; use log::error; use rustpython_ast::{Constant, Expr, ExprContext, ExprKind, Keyword, KeywordData, Stmt, StmtKind}; +use crate::ast::helpers::match_module_member; use crate::ast::types::Range; use crate::autofix::Fix; use crate::check_ast::Checker; @@ -10,11 +11,13 @@ use crate::code_gen::SourceGenerator; use crate::python::identifiers::IDENTIFIER_REGEX; use crate::python::keyword::KWLIST; -/// Return the class name, arguments, and keywords for a `TypedDict` assignment. +/// Return the class name, arguments, keywords and base class for a `TypedDict` +/// assignment. fn match_typed_dict_assign<'a>( + checker: &Checker, targets: &'a [Expr], value: &'a Expr, -) -> Option<(&'a str, &'a [Expr], &'a [Keyword])> { +) -> Option<(&'a str, &'a [Expr], &'a [Keyword], &'a ExprKind)> { if let Some(target) = targets.get(0) { if let ExprKind::Name { id: class_name, .. } = &target.node { if let ExprKind::Call { @@ -23,10 +26,14 @@ fn match_typed_dict_assign<'a>( keywords, } = &value.node { - if let ExprKind::Name { id: func_name, .. } = &func.node { - if func_name == "TypedDict" { - return Some((class_name, args, keywords)); - } + if match_module_member( + func, + "typing", + "TypedDict", + &checker.from_imports, + &checker.import_aliases, + ) { + return Some((class_name, args, keywords, &func.node)); } } } @@ -65,12 +72,13 @@ fn create_pass_stmt() -> Stmt { Stmt::new(Default::default(), Default::default(), StmtKind::Pass) } -/// Generate a `StmtKind:ClassDef` statement bsaed on the provided body and -/// keywords. +/// Generate a `StmtKind:ClassDef` statement based on the provided body, +/// keywords and base class. fn create_class_def_stmt( class_name: &str, body: Vec, total_keyword: Option, + base_class: &ExprKind, ) -> Stmt { let keywords = match total_keyword { Some(keyword) => vec![Keyword::new( @@ -88,10 +96,7 @@ fn create_class_def_stmt( bases: vec![Expr::new( Default::default(), Default::default(), - ExprKind::Name { - id: "TypedDict".to_string(), - ctx: ExprContext::Load, - }, + base_class.clone(), )], keywords, body, @@ -150,11 +155,11 @@ fn get_properties_from_keywords(keywords: &[Keyword]) -> Result> { // The only way to have the `total` keyword is to use the args version, like: // (`TypedDict('name', {'a': int}, total=True)`) -fn get_total_from_only_keyword(keywords: &[Keyword]) -> Option { +fn get_total_from_only_keyword(keywords: &[Keyword]) -> Option<&KeywordData> { match keywords.get(0) { Some(keyword) => match &keyword.node.arg { Some(arg) => match arg.as_str() { - "total" => Some(keyword.node.clone()), + "total" => Some(&keyword.node), _ => None, }, None => None, @@ -171,7 +176,7 @@ fn get_properties_and_total( // dict and keywords. For example, the following is illegal: // MyType = TypedDict('MyType', {'a': int, 'b': str}, a=int, b=str) if let Some(dict) = args.get(1) { - let total = get_total_from_only_keyword(keywords); + let total = get_total_from_only_keyword(keywords).cloned(); match &dict.node { ExprKind::Dict { keys, values } => { Ok((get_properties_from_dict_literal(keys, values)?, total)) @@ -188,15 +193,21 @@ fn get_properties_and_total( } } -/// Generate a `Fix` to convert a `TypedDict` to a functional class. -fn convert_to_functional_class( +/// Generate a `Fix` to convert a `TypedDict` from functional to class. +fn convert_to_class( stmt: &Stmt, class_name: &str, body: Vec, total_keyword: Option, + base_class: &ExprKind, ) -> Result { let mut generator = SourceGenerator::new(); - generator.unparse_stmt(&create_class_def_stmt(class_name, body, total_keyword))?; + generator.unparse_stmt(&create_class_def_stmt( + class_name, + body, + total_keyword, + base_class, + ))?; let content = generator.generate()?; Ok(Fix::replacement( content, @@ -212,7 +223,9 @@ pub fn convert_typed_dict_functional_to_class( targets: &[Expr], value: &Expr, ) { - if let Some((class_name, args, keywords)) = match_typed_dict_assign(targets, value) { + if let Some((class_name, args, keywords, base_class)) = + match_typed_dict_assign(checker, targets, value) + { match get_properties_and_total(args, keywords) { Err(err) => error!("Failed to parse TypedDict: {}", err), Ok((body, total_keyword)) => { @@ -221,7 +234,7 @@ pub fn convert_typed_dict_functional_to_class( Range::from_located(stmt), ); if checker.patch(check.kind.code()) { - match convert_to_functional_class(stmt, class_name, body, total_keyword) { + match convert_to_class(stmt, class_name, body, total_keyword, base_class) { Ok(fix) => check.amend(fix), Err(err) => error!("Failed to convert TypedDict: {}", err), }; diff --git a/src/snapshots/ruff__linter__tests__U013_U013.py.snap b/src/snapshots/ruff__linter__tests__U013_U013.py.snap index d44748c578..858565ebc3 100644 --- a/src/snapshots/ruff__linter__tests__U013_U013.py.snap +++ b/src/snapshots/ruff__linter__tests__U013_U013.py.snap @@ -4,155 +4,172 @@ expression: checks --- - kind: ConvertTypedDictFunctionalToClass location: - row: 4 + row: 5 column: 0 end_location: - row: 4 + row: 5 column: 52 fix: patch: content: "class MyType1(TypedDict):\n a: int\n b: str" location: - row: 4 + row: 5 column: 0 end_location: - row: 4 + row: 5 column: 52 applied: false - kind: ConvertTypedDictFunctionalToClass location: - row: 7 + row: 8 column: 0 end_location: - row: 7 + row: 8 column: 50 fix: patch: content: "class MyType2(TypedDict):\n a: int\n b: str" location: - row: 7 + row: 8 column: 0 end_location: - row: 7 + row: 8 column: 50 applied: false - kind: ConvertTypedDictFunctionalToClass location: - row: 10 + row: 11 column: 0 end_location: - row: 10 + row: 11 column: 44 fix: patch: content: "class MyType3(TypedDict):\n a: int\n b: str" location: - row: 10 + row: 11 column: 0 end_location: - row: 10 + row: 11 column: 44 applied: false - kind: ConvertTypedDictFunctionalToClass location: - row: 13 + row: 14 column: 0 end_location: - row: 13 + row: 14 column: 30 fix: patch: content: "class MyType4(TypedDict):\n pass" location: - row: 13 + row: 14 column: 0 end_location: - row: 13 + row: 14 column: 30 applied: false - kind: ConvertTypedDictFunctionalToClass location: - row: 16 + row: 17 column: 0 end_location: - row: 16 + row: 17 column: 46 fix: patch: content: "class MyType5(TypedDict):\n a: 'hello'" location: - row: 16 + row: 17 column: 0 end_location: - row: 16 + row: 17 column: 46 applied: false - kind: ConvertTypedDictFunctionalToClass location: - row: 17 + row: 18 column: 0 end_location: - row: 17 + row: 18 column: 41 fix: patch: content: "class MyType6(TypedDict):\n a: 'hello'" location: - row: 17 + row: 18 column: 0 end_location: - row: 17 + row: 18 column: 41 applied: false - kind: ConvertTypedDictFunctionalToClass location: - row: 20 + row: 21 column: 0 end_location: - row: 20 + row: 21 column: 56 fix: patch: content: "class MyType7(TypedDict):\n a: NotRequired[dict]" location: - row: 20 + row: 21 column: 0 end_location: - row: 20 + row: 21 column: 56 applied: false - kind: ConvertTypedDictFunctionalToClass location: - row: 23 + row: 24 column: 0 end_location: - row: 23 + row: 24 column: 65 fix: patch: content: "class MyType8(TypedDict, total=False):\n x: int\n y: int" location: - row: 23 + row: 24 column: 0 end_location: - row: 23 + row: 24 column: 65 applied: false - kind: ConvertTypedDictFunctionalToClass location: - row: 29 + row: 30 column: 0 end_location: - row: 29 + row: 30 column: 59 fix: patch: content: "class MyType10(TypedDict):\n key: Literal['value']" location: - row: 29 + row: 30 column: 0 end_location: - row: 29 + row: 30 column: 59 applied: false +- kind: ConvertTypedDictFunctionalToClass + location: + row: 33 + column: 0 + end_location: + row: 33 + column: 53 + fix: + patch: + content: "class MyType11(typing.TypedDict):\n key: int" + location: + row: 33 + column: 0 + end_location: + row: 33 + column: 53 + applied: false