@@ -41,20 +41,20 @@ def get_hints(
4141 file_info = lsp_context .map [path ]
4242
4343 # Process based on whether it's a model or standalone audit
44- if isinstance (file_info , ModelTarget ):
45- # It's a model
46- model = lsp_context .context .get_model (
47- model_or_snapshot = file_info .names [0 ], raise_if_missing = False
48- )
49- if model is None or not isinstance (model , SqlModel ):
50- return []
44+ if not isinstance (file_info , ModelTarget ):
45+ return []
5146
52- query = model .query
53- dialect = model .dialect
54- columns_to_types = model .columns_to_types or {}
55- else :
47+ # It's a model
48+ model = lsp_context .context .get_model (
49+ model_or_snapshot = file_info .names [0 ], raise_if_missing = False
50+ )
51+ if not isinstance (model , SqlModel ):
5652 return []
5753
54+ query = model .query
55+ dialect = model .dialect
56+ columns_to_types = model .columns_to_types or {}
57+
5858 return _get_type_hints_for_model_from_query (
5959 query , dialect , columns_to_types , start_line , end_line
6060 )
@@ -77,32 +77,32 @@ def _get_type_hints_for_model_from_query(
7777
7878 for select in find_all_in_scope (root .expression , exp .Select ):
7979 for select_exp in select .expressions :
80- if not select_exp :
81- continue
82-
8380 if isinstance (select_exp , exp .Alias ):
8481 meta = select_exp .args ["alias" ]._meta
8582 elif isinstance (select_exp , exp .Column ):
8683 meta = select_exp .parts [- 1 ]._meta
8784 else :
8885 continue
8986
90- line = meta .get ("line" ) - 1 # Lines from sqlglot are 1 based
91- col = meta .get ("col" )
92-
93- name = select_exp .alias_or_name
94- if name not in columns_to_types :
87+ if "line" not in meta or "col" not in meta :
9588 continue
9689
90+ line = meta ["line" ]
91+ col = meta ["col" ]
92+
93+ # Lines from sqlglot are 1 based
94+ line -= 1
95+
9796 if line < start_line or line > end_line :
9897 continue
9998
99+ name = select_exp .alias_or_name
100100 data_type = columns_to_types .get (name )
101101
102102 if not data_type or data_type .is_type (exp .DataType .Type .UNKNOWN ):
103103 continue
104104
105- type_label = str ( data_type )
105+ type_label = data_type . sql ( dialect )
106106 hints .append (
107107 types .InlayHint (
108108 label = f"::{ type_label } " ,
0 commit comments