From 112e9d2d8258e72dcf5c21c7d811957f12325af1 Mon Sep 17 00:00:00 2001 From: Rebecca Chen Date: Sun, 15 Dec 2024 04:04:51 -0800 Subject: [PATCH] [ruff_python_ast] Add name and default functions to TypeParam. (#14964) ## Summary This change adds `name` and `default` functions to `TypeParam` to access the corresponding attributes more conveniently. I currently have these as helper functions in code built on top of ruff_python_ast, and they seemed like they might be generally useful. ## Test Plan Ran the checks listed in CONTRIBUTING.md#development. --------- Co-authored-by: Micha Reiser Co-authored-by: Alex Waygood --- crates/ruff_python_ast/src/nodes.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/crates/ruff_python_ast/src/nodes.rs b/crates/ruff_python_ast/src/nodes.rs index c2676da825..4904246eed 100644 --- a/crates/ruff_python_ast/src/nodes.rs +++ b/crates/ruff_python_ast/src/nodes.rs @@ -3386,6 +3386,24 @@ pub enum TypeParam { TypeVarTuple(TypeParamTypeVarTuple), } +impl TypeParam { + pub const fn name(&self) -> &Identifier { + match self { + Self::TypeVar(x) => &x.name, + Self::ParamSpec(x) => &x.name, + Self::TypeVarTuple(x) => &x.name, + } + } + + pub fn default(&self) -> Option<&Expr> { + match self { + Self::TypeVar(x) => x.default.as_deref(), + Self::ParamSpec(x) => x.default.as_deref(), + Self::TypeVarTuple(x) => x.default.as_deref(), + } + } +} + /// See also [TypeVar](https://docs.python.org/3/library/ast.html#ast.TypeVar) #[derive(Clone, Debug, PartialEq)] pub struct TypeParamTypeVar {