Repository: mlflow/mlflow
Pull request: Use Resolver for TYPE_CHECKING Detection in clint linter
<details><summary>&#x1F6E0 DevTools &#x1F6E0</summary>
<p>

[![Open in GitHub Codespaces](https://github.com/codespaces/badge.svg)](https://codespaces.new/Copilot/mlflow/pull/16367?quickstart=1)

#### Install mlflow from this PR

```
# mlflow
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/16367/merge
# mlflow-skinny
pip install git+https://github.com/mlflow/mlflow.git@refs/pull/16367/merge#subdirectory=skinny
```

For Databricks, use the following command:

```
%sh curl -LsSf https://raw.githubusercontent.com/mlflow/mlflow/HEAD/dev/install-skinny.sh | sh -s pull/16367/merge
```

</p>
</details>

## Problem

The clint linter's TYPE_CHECKING detection was hardcoded to only recognize direct usage of the name `TYPE_CHECKING`. This missed valid uses of TYPE_CHECKING with aliases or module imports:

```python
# ✅ Previously worked
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    import some_module

# ❌ Previously missed
from typing import TYPE_CHECKING as TC
if TC:
    import some_module

# ❌ Previously missed  
import typing
if typing.TYPE_CHECKING:
    import some_module
```

## Solution

Enhanced the `visit_If` method in `dev/clint/src/clint/linter.py` to use the existing resolver for proper import resolution:

**Before:**
```python
def visit_If(self, node: ast.If) -> None:
    if isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING":
        self.in_TYPE_CHECKING = True
    self.generic_visit(node)
    self.in_TYPE_CHECKING = False
```

**After:**
```python
def visit_If(self, node: ast.If) -> None:
    # Check if this is a TYPE_CHECKING condition
    # Resolve the test condition to check for TYPE_CHECKING
    if resolved := self.resolver.resolve(node.test):
        # Check if it resolves to typing.TYPE_CHECKING
        if resolved == ["typing", "TYPE_CHECKING"]:
            self.in_TYPE_CHECKING = True

    self.generic_visit(node)
    self.in_TYPE_CHECKING = False
```

## Testing

Validated all documented scenarios now work correctly:

- ✅ Direct import: `from typing import TYPE_CHECKING`
- ✅ Import with alias: `from typing import TYPE_CHECKING as TC`  
- ✅ Module import: `import typing; if typing.TYPE_CHECKING:`
- ✅ Multiple TYPE_CHECKING blocks with different aliases
- ✅ Nested conditions within TYPE_CHECKING blocks
- ✅ Mixed import patterns in the same file

Confirmed no regressions by testing existing mlflow files that use TYPE_CHECKING.

Fixes #16366.

<!-- START COPILOT CODING AGENT TIPS -->
---

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click [here](https://survey.alchemer.com/s3/8343779/Copilot-Coding-agent) to start the survey.

Problem statement:
Use Resolver for TYPE_CHECKING Detection in clint linter
## Problem Statement

Currently, in `dev/clint/src/clint/linter.py`, the detection of `TYPE_CHECKING` usage is hardcoded to only recognize direct usage of the name `TYPE_CHECKING`. This doesn't handle cases where `TYPE_CHECKING` is imported with an alias or from a different module path.

## Motivation

The current implementation misses valid uses of TYPE_CHECKING such as:

- `from typing import TYPE_CHECKING as TC`
- `import typing; if typing.TYPE_CHECKING:`

By using the resolver, we can properly track imports and detect all valid uses of TYPE_CHECKING regardless of how it's imported.

## Solution Overview

Enhance the TYPE_CHECKING detection in the linter to use the resolver for proper import resolution, allowing it to handle various import patterns and aliases.

## Implementation Details

### Files to Modify

1. **`dev/clint/src/clint/linter.py`**
   - Modify the `visit_If` method to use the resolver
   - Update how `in_TYPE_CHECKING` is set

### Step-by-Step Implementation

1. **Update the `visit_If` method** (~lines 692-697):

   Current implementation:

   ```python
   def visit_If(self, node: ast.If) -> None:
       if isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING":
           self.in_TYPE_CHECKING = True
       self.generic_visit(node)
       self.in_TYPE_CHECKING = False
   ```

   New implementation:

   ```python
   def visit_If(self, node: ast.If) -> None:
       # Check if this is a TYPE_CHECKING condition
       # Resolve the test condition to check for TYPE_CHECKING
       if resolved := self.resolver.resolve(node.test):
           # Check if it resolves to typing.TYPE_CHECKING
           if resolved == ["typing", "TYPE_CHECKING"]:
               self.in_TYPE_CHECKING = True

       self.generic_visit(node)
       self.in_TYPE_CHECKING = False
   ```

2. **Handle edge cases**:
   - The resolver already handles various import patterns through `add_import` and `add_import_from`
   - It properly tracks aliases and module paths
   - The `resolve` method will return the fully qualified name

### Example Cases to Handle

1. Direct import:

   ```python
   from typing import TYPE_CHECKING
   if TYPE_CHECKING:  # Should resolve to ["typing", "TYPE_CHECKING"]
   ```

2. Import with alias:

   ```python
   from typing import TYPE_CHECKING as TC
   if TC:  # Should resolve to ["typing", "TYPE_CHECKING"]
   ```

3. Module import:
   ```python
   import typing
   if typing.TYPE_CHECKING:  # Should resolve to ["typing", "TYPE_CHECKING"]
   ```

## Testing

The implementation should be tested with various import patterns to ensure all valid uses of TYPE_CHECKING are properly detected:

1. Test files with different import styles
2. Verify that `in_TYPE_CHECKING` is correctly set within TYPE_CHECKING blocks
3. Ensure that imports are properly tracked when inside TYPE_CHECKING blocks (for the lazy module checking)

Hints from issue discussion:

Commit message: Implement resolver-based TYPE_CHECKING detection in clint linter

Co-authored-by: harupy <17039389+harupy@users.noreply.github.com>

Patch to review:
diff --git a/dev/clint/src/clint/linter.py b/dev/clint/src/clint/linter.py
index b211bde4391d9..80517e6abbf3d 100644
--- a/dev/clint/src/clint/linter.py
+++ b/dev/clint/src/clint/linter.py
@@ -690,8 +690,13 @@ def visit_type_annotation(self, node: ast.AST) -> None:
         self.in_type_annotation = False
 
     def visit_If(self, node: ast.If) -> None:
-        if isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING":
-            self.in_TYPE_CHECKING = True
+        # Check if this is a TYPE_CHECKING condition
+        # Resolve the test condition to check for TYPE_CHECKING
+        if resolved := self.resolver.resolve(node.test):
+            # Check if it resolves to typing.TYPE_CHECKING
+            if resolved == ["typing", "TYPE_CHECKING"]:
+                self.in_TYPE_CHECKING = True
+
         self.generic_visit(node)
         self.in_TYPE_CHECKING = False
 


Code review comments:
diff_hunk: @@ -690,8 +690,13 @@ def visit_type_annotation(self, node: ast.AST) -> None:
         self.in_type_annotation = False
 
     def visit_If(self, node: ast.If) -> None:
-        if isinstance(node.test, ast.Name) and node.test.id == "TYPE_CHECKING":
-            self.in_TYPE_CHECKING = True
+        # Check if this is a TYPE_CHECKING condition
path: dev/clint/src/clint/linter.py
text: @user1:
let's remove unnecessary comments and merge the if statements.

@author:
Removed unnecessary comments and merged if statements as requested. Commit 0836c88